





<!DOCTYPE html>
<html class="writer-html5" lang="zh-CN" >
<head>
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>Bring Your Own Codegen To TVM &mdash; tvm 0.8.dev1982 文档</title>
  

  
  <link rel="stylesheet" href="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/css/bootstrap.min.css" integrity="sha384-Gn5384xqQ1aoWXA+058RXPxPg6fy4IWvTNh0E263XmFcJlSAwiGgFAW/dAiS6JXm" crossorigin="anonymous">
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/gallery.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/css/tlcpack_theme.css" type="text/css" />

  
  
    <link rel="shortcut icon" href="../../_static/tvm-logo-square.png"/>
  

  
  
  
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../" src="../../_static/documentation_options.js"></script>
        <script data-url_root="../../" id="documentation_options" src="../../_static/documentation_options.js"></script>
        <script src="../../_static/jquery.js"></script>
        <script src="../../_static/underscore.js"></script>
        <script src="../../_static/doctools.js"></script>
        <script src="../../_static/translations.js"></script>
    
    <script type="text/javascript" src="../../_static/js/theme.js"></script>

    
    <script type="text/javascript" src="../../_static/js/tlcpack_theme.js"></script>
    <link rel="index" title="索引" href="../../genindex.html" />
    <link rel="search" title="搜索" href="../../search.html" />
    <link rel="next" title="Python Target Parametrization" href="pytest_target_parametrization.html" />
    <link rel="prev" title="Adding a Compiler Pass to Relay" href="relay_add_pass.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    
<header class="header">
    <div class="innercontainer">
      <div class="headerInner d-flex justify-content-between align-items-center">
          <div class="headerLogo">
               <a href="https://tvm.apache.org/"><img src=https://tvm.apache.org/assets/images/logo.svg alt="logo"></a>
          </div>

          <div id="headMenu" class="headerNav">
            <button type="button" id="closeHeadMenu" class="navCloseBtn"><img src="../../_static/img/close-icon.svg" alt="Close"></button>
             <ul class="nav">
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/community>Community</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/download>Download</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/vta>VTA</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/blog>Blog</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/docs>Docs</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvmconf.org>Conference</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://github.com/apache/tvm/>Github</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvmchinese.github.io/declaration_zh_CN.html>About-Translators</a>
                </li>
             </ul>
               <div class="responsivetlcdropdown">
                 <button type="button" class="btn-link">
                   ASF
                 </button>
                 <ul>
                     <li>
                       <a href=https://apache.org/>Apache Homepage</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/licenses/>License</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/sponsorship.html>Sponsorship</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/security/>Security</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/thanks.html>Thanks</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/events/current-event>Events</a>
                     </li>
                     <li>
                       <a href=https://www.zhihu.com/column/c_1429578595417563136>Zhihu</a>
                     </li>
                 </ul>
               </div>
          </div>
            <div class="responsiveMenuIcon">
              <button type="button" id="menuBtn" class="btn-menu"><img src="../../_static/img/menu-icon.svg" alt="Menu Icon"></button>
            </div>

            <div class="tlcDropdown">
              <div class="dropdown">
                <button type="button" class="btn-link dropdown-toggle" data-toggle="dropdown" aria-haspopup="true" aria-expanded="false">
                  ASF
                </button>
                <div class="dropdown-menu dropdown-menu-right">
                  <ul>
                     <li>
                       <a href=https://apache.org/>Apache Homepage</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/licenses/>License</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/sponsorship.html>Sponsorship</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/security/>Security</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/thanks.html>Thanks</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/events/current-event>Events</a>
                     </li>
                     <li>
                       <a href=https://www.zhihu.com/column/c_1429578595417563136>Zhihu</a>
                     </li>
                  </ul>
                </div>
              </div>
          </div>
       </div>
    </div>
 </header>
 
    <nav data-toggle="wy-nav-shift" class="wy-nav-side fixed">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../index.html">
          

          
            
            <img src="../../_static/tvm-logo-small.png" class="logo" alt="Logo"/>
          
          </a>

          
            
            
                <div class="version">
                  0.8.dev1982
                </div>
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption" role="heading"><span class="caption-text">如何开始</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../install/index.html">安装 TVM</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../contribute/index.html">贡献者指南</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">用户引导</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../tutorial/index.html">User Tutorial</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../how_to/index.html">How To Guides</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">开发者引导</span></p>
<ul class="current">
<li class="toctree-l1"><a class="reference internal" href="../tutorial/index.html">Developer Tutorial</a></li>
<li class="toctree-l1 current"><a class="reference internal" href="how_to.html">开发者指南</a><ul class="current">
<li class="toctree-l2"><a class="reference internal" href="relay_add_op.html">向 Realy 中添加一个算子</a></li>
<li class="toctree-l2"><a class="reference internal" href="relay_add_pass.html">Adding a Compiler Pass to Relay</a></li>
<li class="toctree-l2 current"><a class="current reference internal" href="#">Bring Your Own Codegen To TVM</a><ul>
<li class="toctree-l3"><a class="reference internal" href="#implement-a-c-codegen">Implement a C Codegen</a><ul>
<li class="toctree-l4"><a class="reference internal" href="#implement-codegenc">Implement CodegenC</a></li>
<li class="toctree-l4"><a class="reference internal" href="#implement-csourcecodegen">Implement CSourceCodegen</a></li>
<li class="toctree-l4"><a class="reference internal" href="#register-your-codegen">Register Your Codegen</a></li>
</ul>
</li>
<li class="toctree-l3"><a class="reference internal" href="#implement-a-codegen-for-your-representation">Implement a Codegen for Your Representation</a><ul>
<li class="toctree-l4"><a class="reference internal" href="#implement-examplejsoncodegen">Implement ExampleJsonCodeGen</a></li>
<li class="toctree-l4"><a class="reference internal" href="#implement-a-customized-runtime">Implement a Customized Runtime</a></li>
</ul>
</li>
<li class="toctree-l3"><a class="reference internal" href="#summary">总结</a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="pytest_target_parametrization.html">Python Target Parametrization</a></li>
</ul>
</li>
</ul>
<p class="caption" role="heading"><span class="caption-text">架构指南</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../arch/index.html">Design and Architecture</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">主题引导</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../topic/microtvm/index.html">microTVM：裸机使用TVM</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../topic/vta/index.html">VTA: Versatile Tensor Accelerator</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">参考指南</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../reference/langref/index.html">语言参考</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/api/python/index.html">Python API</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/api/links.html">Other APIs</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/publications.html">Publications</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../genindex.html">索引</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">
      
      <nav class="wy-nav-top" aria-label="top navigation" data-toggle="wy-nav-top">
        
            <div class="togglemenu">

            </div>
            <div class="nav-content">
              <!-- tvm -->
              Table of content
            </div>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        

          




















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../index.html">Docs</a> <span class="br-arrow">></span></li>
        
          <li><a href="how_to.html">开发者指南</a> <span class="br-arrow">></span></li>
        
      <li>Bring Your Own Codegen To TVM</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
            
            <a href="../../_sources/dev/how_to/relay_bring_your_own_codegen.rst.txt" rel="nofollow"> <img src="../../_static//img/source.svg" alt="viewsource"/></a>
          
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <div class="section" id="bring-your-own-codegen-to-tvm">
<span id="relay-bring-your-own-codegen"></span><h1>Bring Your Own Codegen To TVM<a class="headerlink" href="#bring-your-own-codegen-to-tvm" title="永久链接至标题">¶</a></h1>
<p>As the number of hardware devices targeted by deep learning workloads keeps increasing, the required knowledge for users to achieve high performance on various devices keeps increasing as well. To free data scientists from worrying about the performance when developing a new model, hardware backend providers either provide libraries such as MKLDNN or cuDNN with many commonly used deep learning operators, or provide frameworks such as TensorRT to let users describe their models in a certain way to achieve high performance. However, users have to learn a new programming interface when they attempt to work on a new library or device. As a result, the demand for a unified programming interface becomes more and more important to 1) let all users and hardware backend providers stand on the same page, and 2) provide a feasible solution to allow specialized hardware or library to only support widely used operators with extremely high performance, but fallback unsupported operators to general devices like CPU/GPU.</p>
<p>In this developer guide, we demonstrate how you, as a hardware backend provider, can easily implement your own codegen and register it as a Relay backend compiler to support your hardware device/library. This guide covers two types of codegen based on different graph representations you need:</p>
<p><strong>1. You want to generate C code.</strong></p>
<p>If your hardware already has a well-optimized C/C++ library, such as Intel CBLAS/MKL to CPU and NVIDIA CUBLAS to GPU, then this is what you are looking for. Fortunately, C source code module is fully compatible with TVM runtime module, which means the generated code could be compiled by any C/C++ compiler with proper compilation flags, so the only task you have is to implement a codegen that generates C code for subgraphs and a C source module to integrate into TVM runtime module. We will demonstrate how to implement a C code generator for your hardware in the following section.</p>
<p><strong>2. You want to generate any other graph representations.</strong></p>
<p>Your hardware may require other forms of graph representation, such as JSON. In this case, you need to implement not only a codegen but also a customized TVM runtime module to let TVM runtime know how this graph representation should be executed. If you already have a complete graph execution engine for your hardware, such as TensorRT for GPU, then this is a solution you can consider.</p>
<p>After you finish the codegen and runtime, you can then let your customers annotate their models with your customized tag to make use of them. The tutorial for end-users to annotate and launch a specific codegen is <strong>here (TBA)</strong>.</p>
<div class="section" id="implement-a-c-codegen">
<h2>Implement a C Codegen<a class="headerlink" href="#implement-a-c-codegen" title="永久链接至标题">¶</a></h2>
<p>In this part, we demonstrate how to implement a codegen that generates C code with pre-implemented operator functions. To simplify, our example codegen does not depend on third-party libraries. Instead, we manually implement two macros in C:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="cp">#define CSOURCE_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_)         \</span>
<span class="cp">    extern &quot;C&quot; void p_ID_(float* a, float* b, float* out) { \</span>
<span class="cp">        for (int64_t i = 0; i &lt; p_DIM1_; ++i) {             \</span>
<span class="cp">            out[i] = a[i] p_OP_ b[i];                       \</span>
<span class="cp">        }                                                   \</span>
<span class="cp">    }</span>

<span class="cp">#define CSOURCE_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_)  \</span>
<span class="cp">    extern &quot;C&quot; void p_ID_(float* a, float* b, float* out) {   \</span>
<span class="cp">        for (int64_t i = 0; i &lt; p_DIM1_; ++i) {               \</span>
<span class="cp">            for (int64_t j = 0; j &lt; p_DIM2_; ++j) {           \</span>
<span class="cp">                int64_t k = i * p_DIM2_ + j;                  \</span>
<span class="cp">                out[k] = a[k] p_OP_ b[k];                     \</span>
<span class="cp">            }                                                 \</span>
<span class="cp">        }                                                     \</span>
<span class="cp">    }</span>
</pre></div>
</div>
<p>With the two macros, we can generate binary operators for 1-D and 2-D tensors. For example, given a subgraph as follows. Assuming all inputs are 2-D tensors with shape (10, 10).</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">c_compiler_input0</span>
       <span class="o">|</span>
      <span class="n">add</span> <span class="o">&lt;--</span> <span class="n">c_compiler_input1</span>
       <span class="o">|</span>
    <span class="n">subtract</span> <span class="o">&lt;--</span> <span class="n">c_compiler_input2</span>
       <span class="o">|</span>
    <span class="n">multiply</span> <span class="o">&lt;--</span> <span class="n">c_compiler_input3</span>
       <span class="o">|</span>
      <span class="n">out</span>
</pre></div>
</div>
<p>Our goal is to generate the following compilable code to execute the subgraph:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/c_runtime_api.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/packed_func.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;dlpack/dlpack.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;cstdint&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;cstring&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;iostream&gt;</span><span class="cp"></span>

<span class="cp">#define GCC_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_)           \</span>
<span class="cp">  extern &quot;C&quot; void p_ID_(float* a, float* b, float* out) { \</span>
<span class="cp">    for (int64_t i = 0; i &lt; p_DIM1_; ++i) {               \</span>
<span class="cp">      out[i] = a[i] p_OP_ b[i];                           \</span>
<span class="cp">    }                                                     \</span>
<span class="cp">  }</span>

<span class="cp">#define GCC_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_)  \</span>
<span class="cp">  extern &quot;C&quot; void p_ID_(float* a, float* b, float* out) { \</span>
<span class="cp">    for (int64_t i = 0; i &lt; p_DIM1_; ++i) {               \</span>
<span class="cp">      for (int64_t j = 0; j &lt; p_DIM2_; ++j) {             \</span>
<span class="cp">        int64_t k = i * p_DIM2_ + j;                      \</span>
<span class="cp">        out[k] = a[k] p_OP_ b[k];                         \</span>
<span class="cp">      }                                                   \</span>
<span class="cp">    }                                                     \</span>
<span class="cp">  }</span>

<span class="c1">// Note 1</span>
<span class="n">GCC_BINARY_OP_2D</span><span class="p">(</span><span class="n">gcc_0_0</span><span class="p">,</span> <span class="o">*</span><span class="p">,</span> <span class="mi">10</span><span class="p">,</span> <span class="mi">10</span><span class="p">);</span>
<span class="n">GCC_BINARY_OP_2D</span><span class="p">(</span><span class="n">gcc_0_1</span><span class="p">,</span> <span class="o">-</span><span class="p">,</span> <span class="mi">10</span><span class="p">,</span> <span class="mi">10</span><span class="p">);</span>
<span class="n">GCC_BINARY_OP_2D</span><span class="p">(</span><span class="n">gcc_0_2</span><span class="p">,</span> <span class="o">+</span><span class="p">,</span> <span class="mi">10</span><span class="p">,</span> <span class="mi">10</span><span class="p">);</span>

<span class="c1">// Note 2</span>
<span class="k">extern</span> <span class="s">&quot;C&quot;</span> <span class="kt">void</span> <span class="n">gcc_0_</span><span class="p">(</span><span class="kt">float</span><span class="o">*</span> <span class="n">gcc_input0</span><span class="p">,</span> <span class="kt">float</span><span class="o">*</span> <span class="n">gcc_input1</span><span class="p">,</span>
                       <span class="kt">float</span><span class="o">*</span> <span class="n">gcc_input2</span><span class="p">,</span> <span class="kt">float</span><span class="o">*</span> <span class="n">gcc_input3</span><span class="p">,</span> <span class="kt">float</span><span class="o">*</span> <span class="n">out</span><span class="p">)</span> <span class="p">{</span>
  <span class="kt">float</span><span class="o">*</span> <span class="n">buf_0</span> <span class="o">=</span> <span class="p">(</span><span class="kt">float</span><span class="o">*</span><span class="p">)</span><span class="n">malloc</span><span class="p">(</span><span class="mi">4</span> <span class="o">*</span> <span class="mi">100</span><span class="p">);</span>
  <span class="kt">float</span><span class="o">*</span> <span class="n">buf_1</span> <span class="o">=</span> <span class="p">(</span><span class="kt">float</span><span class="o">*</span><span class="p">)</span><span class="n">malloc</span><span class="p">(</span><span class="mi">4</span> <span class="o">*</span> <span class="mi">100</span><span class="p">);</span>
  <span class="n">gcc_0_2</span><span class="p">(</span><span class="n">gcc_input0</span><span class="p">,</span> <span class="n">gcc_input1</span><span class="p">,</span> <span class="n">buf_0</span><span class="p">);</span>
  <span class="n">gcc_0_1</span><span class="p">(</span><span class="n">buf_0</span><span class="p">,</span> <span class="n">gcc_input2</span><span class="p">,</span> <span class="n">buf_1</span><span class="p">);</span>
  <span class="n">gcc_0_0</span><span class="p">(</span><span class="n">buf_1</span><span class="p">,</span> <span class="n">gcc_input3</span><span class="p">,</span> <span class="n">out</span><span class="p">);</span>
  <span class="n">free</span><span class="p">(</span><span class="n">buf_0</span><span class="p">);</span>
  <span class="n">free</span><span class="p">(</span><span class="n">buf_1</span><span class="p">);</span>
<span class="p">}</span>

<span class="c1">// Note 3</span>
<span class="k">extern</span> <span class="s">&quot;C&quot;</span> <span class="kt">int</span> <span class="n">gcc_0_wrapper</span><span class="p">(</span><span class="n">DLTensor</span><span class="o">*</span> <span class="n">arg0</span><span class="p">,</span> <span class="n">DLTensor</span><span class="o">*</span> <span class="n">arg1</span><span class="p">,</span> <span class="n">DLTensor</span><span class="o">*</span> <span class="n">arg2</span><span class="p">,</span>
                             <span class="n">DLTensor</span><span class="o">*</span> <span class="n">arg3</span><span class="p">,</span> <span class="n">DLTensor</span><span class="o">*</span> <span class="n">out</span><span class="p">)</span> <span class="p">{</span>
  <span class="n">gcc_0_</span><span class="p">(</span><span class="k">static_cast</span><span class="o">&lt;</span><span class="kt">float</span><span class="o">*&gt;</span><span class="p">(</span><span class="n">arg0</span><span class="o">-&gt;</span><span class="n">data</span><span class="p">),</span> <span class="k">static_cast</span><span class="o">&lt;</span><span class="kt">float</span><span class="o">*&gt;</span><span class="p">(</span><span class="n">arg1</span><span class="o">-&gt;</span><span class="n">data</span><span class="p">),</span>
         <span class="k">static_cast</span><span class="o">&lt;</span><span class="kt">float</span><span class="o">*&gt;</span><span class="p">(</span><span class="n">arg2</span><span class="o">-&gt;</span><span class="n">data</span><span class="p">),</span> <span class="k">static_cast</span><span class="o">&lt;</span><span class="kt">float</span><span class="o">*&gt;</span><span class="p">(</span><span class="n">arg3</span><span class="o">-&gt;</span><span class="n">data</span><span class="p">),</span>
         <span class="k">static_cast</span><span class="o">&lt;</span><span class="kt">float</span><span class="o">*&gt;</span><span class="p">(</span><span class="n">out</span><span class="o">-&gt;</span><span class="n">data</span><span class="p">));</span>
  <span class="k">return</span> <span class="mi">0</span><span class="p">;</span>
<span class="p">}</span>
<span class="n">TVM_DLL_EXPORT_TYPED_FUNC</span><span class="p">(</span><span class="n">gcc_0</span><span class="p">,</span> <span class="n">gcc_0_wrapper</span><span class="p">);</span>
</pre></div>
</div>
<p>Here we highlight the notes marked in the above code:</p>
<ul class="simple">
<li><p><strong>Note 1</strong> is the function implementation for the three nodes in the subgraph.</p></li>
<li><p><strong>Note 2</strong> is a function to execute the subgraph by allocating intermediate buffers and invoking corresponding functions.</p></li>
<li><p><strong>Note 3</strong> is a TVM runtime compatible wrapper function. It accepts a list of input tensors and one output tensor (the last argument), casts them to the right data type, and invokes the subgraph function described in Note 2. In addition, <code class="docutils literal notranslate"><span class="pre">TVM_DLL_EXPORT_TYPED_FUNC</span></code> is a TVM macro that generates another function <code class="docutils literal notranslate"><span class="pre">gcc_0</span></code> with unified the function arguments by packing all tensors to <code class="docutils literal notranslate"><span class="pre">TVMArgs</span></code>. As a result, the TVM runtime can directly invoke <code class="docutils literal notranslate"><span class="pre">gcc_0</span></code> to execute the subgraph without additional efforts. With the above code generated, TVM is able to compile it along with the rest parts of the graph and export a single library for deployment.</p></li>
</ul>
<p>In the rest of this section, we will implement a codegen step-by-step to generate the above code. Your own codegen has to be located at <code class="docutils literal notranslate"><span class="pre">src/relay/backend/contrib/&lt;your-codegen-name&gt;/</span></code>. In our example, we name our codegen “codegen_c” and put it under <a class="reference external" href="https://github.com/apache/tvm/blob/main/src/relay/backend/contrib/codegen_c/codegen.cc">/src/relay/backend/contrib/codegen_c/</a>. Feel free to check this file for a complete implementation.</p>
<p>Specifically, we are going to implement two classes in this file and here is their relationship:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span>                     <span class="n">subgraph</span>                                <span class="n">subgraph</span>
<span class="n">TVM</span> <span class="n">backend</span> <span class="o">-----------------------------&gt;</span> <span class="n">CSourceCodegen</span> <span class="o">-------------&gt;</span> <span class="n">CodegenC</span>
       <span class="o">^</span>                                       <span class="o">|</span>    <span class="o">^</span>                       <span class="o">|</span>
       <span class="o">|</span>                                       <span class="o">|</span>    <span class="o">|</span>                       <span class="o">|</span>
       <span class="o">----------------------------------------</span>      <span class="o">------------------------</span>
          <span class="n">generated</span> <span class="n">C</span> <span class="n">source</span> <span class="n">runtime</span> <span class="n">module</span>              <span class="n">generated</span> <span class="n">C</span> <span class="n">code</span>
</pre></div>
</div>
<p>When TVM backend finds a function (subgraph) in a Relay graph is annotated with the registered compiler tag (<code class="docutils literal notranslate"><span class="pre">ccompiler</span></code> in this example), TVM backend invokes <code class="docutils literal notranslate"><span class="pre">CSourceCodegen</span></code> and passes the subgraph. <code class="docutils literal notranslate"><span class="pre">CSourceCodegen</span></code>’s member function <code class="docutils literal notranslate"><span class="pre">CreateCSourceModule</span></code> will 1) generate C code for the subgraph, and 2) wrap the generated C code to a C source runtime module for TVM backend to compile and deploy. In particular, the C code generation is transparent to the <code class="docutils literal notranslate"><span class="pre">CodegenC</span></code> class because it provides many useful utilities to ease the code generation implementation. The following sections will implement these two classes in the bottom-up order.</p>
<div class="section" id="implement-codegenc">
<h3>Implement CodegenC<a class="headerlink" href="#implement-codegenc" title="永久链接至标题">¶</a></h3>
<p>In <code class="docutils literal notranslate"><span class="pre">src/relay/backend/contrib/codegen_c/codegen.cc</span></code>, we first create a codegen class skeleton under the namespace of <code class="docutils literal notranslate"><span class="pre">tvm.relay.contrib</span></code>:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="cp">#include</span> <span class="cpf">&lt;tvm/relay/expr_functor.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/relay/transform.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/relay/type.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/module.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/object.h&gt;</span><span class="cp"></span>

<span class="cp">#include</span> <span class="cpf">&lt;fstream&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;sstream&gt;</span><span class="cp"></span>

<span class="cp">#include</span> <span class="cpf">&quot;codegen_c.h&quot;</span><span class="cp"></span>

<span class="k">namespace</span> <span class="n">tvm</span> <span class="p">{</span>
<span class="k">namespace</span> <span class="n">relay</span> <span class="p">{</span>
<span class="k">namespace</span> <span class="n">contrib</span> <span class="p">{</span>

<span class="k">class</span> <span class="nc">CodegenC</span> <span class="o">:</span> <span class="k">public</span> <span class="n">ExprVisitor</span><span class="p">,</span> <span class="k">public</span> <span class="n">CodegenCBase</span> <span class="p">{</span>
  <span class="k">public</span><span class="o">:</span>
    <span class="k">explicit</span> <span class="n">CodegenC</span><span class="p">(</span><span class="k">const</span> <span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&amp;</span> <span class="n">id</span><span class="p">)</span> <span class="p">{</span> <span class="k">this</span><span class="o">-&gt;</span><span class="n">ext_func_id_</span> <span class="o">=</span> <span class="n">id</span><span class="p">;</span> <span class="p">}</span>

    <span class="kt">void</span> <span class="n">VisitExpr_</span><span class="p">(</span><span class="k">const</span> <span class="n">VarNode</span><span class="o">*</span> <span class="n">node</span><span class="p">)</span> <span class="p">{</span> <span class="p">;</span> <span class="p">}</span>
    <span class="kt">void</span> <span class="n">VisitExpr_</span><span class="p">(</span><span class="k">const</span> <span class="n">CallNode</span><span class="o">*</span> <span class="n">call</span><span class="p">)</span> <span class="k">final</span> <span class="p">{</span> <span class="p">;</span> <span class="p">}</span>
    <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">JIT</span><span class="p">()</span> <span class="p">{</span> <span class="p">;</span> <span class="p">}</span>

  <span class="k">private</span><span class="o">:</span>
    <span class="cm">/*! \brief The function id that represents a C source function. */</span>
    <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">ext_func_id_</span> <span class="o">=</span> <span class="s">&quot;&quot;</span><span class="p">;</span>
    <span class="cm">/*! \brief The index of a wrapped C function. */</span>
    <span class="kt">int</span> <span class="n">func_idx</span> <span class="o">=</span> <span class="mi">0</span><span class="p">;</span>
    <span class="cm">/*! \brief The index of allocated buffers. */</span>
    <span class="kt">int</span> <span class="n">buf_idx_</span> <span class="o">=</span> <span class="mi">0</span><span class="p">;</span>
    <span class="cm">/*! \brief The arguments of a C compiler compatible function. */</span>
    <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&gt;</span> <span class="n">ext_func_args_</span><span class="p">;</span>
    <span class="cm">/*! \brief The statements of a C compiler compatible function. */</span>
    <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&gt;</span> <span class="n">ext_func_body</span><span class="p">;</span>
    <span class="cm">/*! \brief The declaration statements of a C compiler compatible function. */</span>
    <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&gt;</span> <span class="n">func_decl_</span><span class="p">;</span>
    <span class="cm">/*! \brief The declaration statements of buffers. */</span>
    <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&gt;</span> <span class="n">buf_decl_</span><span class="p">;</span>
    <span class="cm">/*! \brief The name and index pairs for output. */</span>
    <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="n">std</span><span class="o">::</span><span class="n">pair</span><span class="o">&lt;</span><span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="p">,</span> <span class="kt">int</span><span class="o">&gt;&gt;</span> <span class="n">out_</span><span class="p">;</span>
<span class="p">}</span>
</pre></div>
</div>
<p>The <code class="docutils literal notranslate"><span class="pre">CodegenC</span></code> class inherits two classes: <code class="docutils literal notranslate"><span class="pre">ExprVisitor</span></code> provides abilities to traverse subgraphs and collects the required information and generate subgraph functions such as <code class="docutils literal notranslate"><span class="pre">gcc_0_</span></code>; <code class="docutils literal notranslate"><span class="pre">CodegenCBase</span></code> provides abilities and utilities to generate wrapper functions such as <code class="docutils literal notranslate"><span class="pre">gcc_0</span></code> in the above example. As can be seen, we only need to implement three functions in this codegen class to make it work.</p>
<div class="section" id="code-generation-for-operators">
<h4>Code Generation for Operators<a class="headerlink" href="#code-generation-for-operators" title="永久链接至标题">¶</a></h4>
<p>We first implement <code class="docutils literal notranslate"><span class="pre">VisitExpr_(const</span> <span class="pre">CallNode*</span> <span class="pre">call)</span></code>. This function visits all call nodes when traversing the subgraph. Each call node contains an operator that we want to offload to your hardware. As a result, we need to generate the corresponding C code with correct operators in topological order. We implement this function step-by-step as follows.</p>
<p><strong>1. Generate the function declaration</strong></p>
<p>Example Result: <code class="docutils literal notranslate"><span class="pre">GCC_BINARY_OP_2D(gcc_0_0,</span> <span class="pre">*,</span> <span class="pre">10,</span> <span class="pre">10);</span></code></p>
<p>To generate the function declaration, as shown above, we need 1) a function name (e.g., <code class="docutils literal notranslate"><span class="pre">gcc_0_0</span></code>), 2) the type of operator (e.g., <code class="docutils literal notranslate"><span class="pre">*</span></code>), and 3) the input tensor shape (e.g., <code class="docutils literal notranslate"><span class="pre">(10,</span> <span class="pre">10)</span></code>). Fortunately, this information can be obtained easily from <code class="docutils literal notranslate"><span class="pre">CallNode</span></code>:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">std</span><span class="o">::</span><span class="n">ostringstream</span> <span class="n">macro_stream</span><span class="p">;</span>
<span class="n">std</span><span class="o">::</span><span class="n">ostringstream</span> <span class="n">decl_stream</span><span class="p">;</span>
<span class="n">std</span><span class="o">::</span><span class="n">ostringstream</span> <span class="n">buf_stream</span><span class="p">;</span>

<span class="c1">// Generate a unique function name you like.</span>
<span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">func_name</span> <span class="o">=</span> <span class="n">ext_func_id_</span> <span class="o">+</span> <span class="s">&quot;_&quot;</span> <span class="o">+</span> <span class="n">std</span><span class="o">::</span><span class="n">to_string</span><span class="p">(</span><span class="n">func_idx</span><span class="o">++</span><span class="p">);</span>

<span class="c1">// Make function declaration string.</span>
<span class="n">macro_stream</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;CSOURCE_BINARY_OP_&quot;</span> <span class="o">&lt;&lt;</span> <span class="n">call</span><span class="o">-&gt;</span><span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">()</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;D(&quot;</span> <span class="o">&lt;&lt;</span> <span class="n">func_name</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;, &quot;</span><span class="p">;</span>

<span class="c1">// Check the operator type.</span>
<span class="k">if</span> <span class="p">(</span><span class="n">IsOp</span><span class="p">(</span><span class="n">call</span><span class="p">,</span> <span class="s">&quot;add&quot;</span><span class="p">))</span> <span class="p">{</span>
  <span class="n">macro_stream</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;+&quot;</span><span class="p">;</span>
<span class="p">}</span> <span class="k">else</span> <span class="k">if</span> <span class="p">(</span><span class="n">IsOp</span><span class="p">(</span><span class="n">call</span><span class="p">,</span> <span class="s">&quot;subtract&quot;</span><span class="p">))</span> <span class="p">{</span>
  <span class="n">macro_stream</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;-&quot;</span><span class="p">;</span>
<span class="p">}</span> <span class="k">else</span> <span class="k">if</span> <span class="p">(</span><span class="n">IsOp</span><span class="p">(</span><span class="n">call</span><span class="p">,</span> <span class="s">&quot;multiply&quot;</span><span class="p">))</span> <span class="p">{</span>
  <span class="n">macro_stream</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;*&quot;</span><span class="p">;</span>
<span class="p">}</span> <span class="k">else</span> <span class="p">{</span>
  <span class="n">LOG</span><span class="p">(</span><span class="n">FATAL</span><span class="p">)</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;Unrecognized op&quot;</span><span class="p">;</span>
<span class="p">}</span>

<span class="c1">// Extract the input tensor shape.</span>
<span class="k">auto</span> <span class="n">in_shape</span> <span class="o">=</span> <span class="n">GetShape</span><span class="p">(</span><span class="n">call</span><span class="o">-&gt;</span><span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">-&gt;</span><span class="n">checked_type</span><span class="p">());</span>
<span class="k">for</span> <span class="p">(</span><span class="kt">size_t</span> <span class="n">i</span> <span class="o">=</span> <span class="mi">0</span><span class="p">;</span> <span class="n">i</span> <span class="o">&lt;</span> <span class="n">in_shape</span><span class="p">.</span><span class="n">size</span><span class="p">();</span> <span class="o">++</span><span class="n">i</span><span class="p">)</span> <span class="p">{</span>
  <span class="n">macro_stream</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;, &quot;</span> <span class="o">&lt;&lt;</span> <span class="n">in_shape</span><span class="p">[</span><span class="n">i</span><span class="p">];</span>
<span class="p">}</span>
<span class="n">macro_stream</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;);&quot;</span><span class="p">;</span>
<span class="n">func_decl_</span><span class="p">.</span><span class="n">push_back</span><span class="p">(</span><span class="n">macro_stream</span><span class="p">.</span><span class="n">str</span><span class="p">());</span>
</pre></div>
</div>
<p>As can be seen, we push the generated code to class member variables <code class="docutils literal notranslate"><span class="pre">func_decl_</span></code>. It means after we finish traversing the entire subgraph, we have collected all required function declarations and the only thing we need to do is having them compiled by GCC. The rest implementation of <code class="docutils literal notranslate"><span class="pre">VisitExpr_(const</span> <span class="pre">CallNode*</span> <span class="pre">call)</span></code> also follow this concept.</p>
<p><strong>2. Generate the function call</strong></p>
<p>Example Result: <code class="docutils literal notranslate"><span class="pre">gcc_0_0(buf_1,</span> <span class="pre">gcc_input3,</span> <span class="pre">out);</span></code></p>
<p>After generating the function declaration, we need to generate a function call with proper inputs and outputs. To know which inputs or buffers we should put when calling this function, we have to visit its arguments:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="kt">bool</span> <span class="n">first</span> <span class="o">=</span> <span class="nb">true</span><span class="p">;</span>
<span class="n">decl_stream</span> <span class="o">&lt;&lt;</span> <span class="n">func_name</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;(&quot;</span><span class="p">;</span>
<span class="k">for</span> <span class="p">(</span><span class="kt">size_t</span> <span class="n">i</span> <span class="o">=</span> <span class="mi">0</span><span class="p">;</span> <span class="n">i</span> <span class="o">&lt;</span> <span class="n">call</span><span class="o">-&gt;</span><span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">();</span> <span class="o">++</span><span class="n">i</span><span class="p">)</span> <span class="p">{</span>
  <span class="n">VisitExpr</span><span class="p">(</span><span class="n">call</span><span class="o">-&gt;</span><span class="n">args</span><span class="p">[</span><span class="n">i</span><span class="p">]);</span> <span class="c1">// Note 1</span>
  <span class="k">for</span> <span class="p">(</span><span class="k">auto</span> <span class="nl">out</span> <span class="p">:</span> <span class="n">out_</span><span class="p">)</span> <span class="p">{</span>
    <span class="k">if</span> <span class="p">(</span><span class="o">!</span><span class="n">first</span><span class="p">)</span> <span class="p">{</span>
      <span class="n">decl_stream</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;, &quot;</span><span class="p">;</span>
    <span class="p">}</span>
    <span class="n">first</span> <span class="o">=</span> <span class="nb">false</span><span class="p">;</span>
    <span class="n">decl_stream</span> <span class="o">&lt;&lt;</span> <span class="n">out</span><span class="p">.</span><span class="n">first</span><span class="p">;</span>
  <span class="p">}</span>
<span class="p">}</span>
<span class="c1">// Note 2</span>
</pre></div>
</div>
<p>Again, we want to highlight the notes in the above code:</p>
<p><strong>Note 1</strong>: <code class="docutils literal notranslate"><span class="pre">VisitExpr(call-&gt;args[i])</span></code> is a recursive call to visit arguments of the current function. An argument could be an output of another node or an input tensor. In our example implementation, we make sure every node updates a class variable <code class="docutils literal notranslate"><span class="pre">out_</span></code> before leaving the visitor. Here is an illustration:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span>  <span class="n">arg_node</span>                 <span class="n">arg_node</span> <span class="o">&lt;-</span> <span class="n">Visit</span> <span class="n">arg</span> <span class="p">(</span><span class="n">Note</span> <span class="mi">1</span><span class="p">)</span>       <span class="n">arg_node</span>
     <span class="o">|</span>                        <span class="o">|</span>                                    <span class="o">|</span>
 <span class="n">curr_node</span> <span class="o">&lt;-</span> <span class="n">Process</span>      <span class="n">curr_node</span>                            <span class="n">curr_node</span> <span class="o">&lt;-</span> <span class="n">Put</span> <span class="s2">&quot;buf_0&quot;</span> <span class="k">as</span> <span class="n">an</span> <span class="nb">input</span> <span class="n">buffer</span>

<span class="p">(</span><span class="n">a</span><span class="p">)</span> <span class="n">out_</span> <span class="o">=</span> <span class="p">{}</span>            <span class="p">(</span><span class="n">b</span><span class="p">)</span> <span class="n">out_</span> <span class="o">=</span> <span class="p">{}</span>                   <span class="p">(</span><span class="n">c</span><span class="p">)</span> <span class="n">out_</span> <span class="o">=</span> <span class="p">{(</span><span class="s2">&quot;buf_0&quot;</span><span class="p">,</span> <span class="mi">20</span><span class="p">)}</span>
</pre></div>
</div>
<p>We can see in the above figure, class variable <code class="docutils literal notranslate"><span class="pre">out_</span></code> is empty before visiting the argument node, and it was filled with the output buffer name and size of <code class="docutils literal notranslate"><span class="pre">arg_node</span></code>. As a result, when we finished visiting the argument node, we know the proper input buffer we should put by looking at <code class="docutils literal notranslate"><span class="pre">out_</span></code>. You will find out how we update <code class="docutils literal notranslate"><span class="pre">out_</span></code> at the end of this section as well as the next section.</p>
<p><strong>Note 2</strong>: You may notice that we did not close the function call string in this step. The current function call string looks like: <code class="docutils literal notranslate"><span class="pre">gcc_0_0(buf_1,</span> <span class="pre">gcc_input3</span></code>. This is because we have not put the last argument (i.e., the output) to this call. The output of a function call could be either an allocated temporary buffer or the subgraph output tensor. For simplify, in this example, we allocate an output buffer for every call node (next step) and copy the result in the very last buffer to the output tensor.</p>
<p><strong>3. Generate the output buffer</strong></p>
<p>Example Result: <code class="docutils literal notranslate"><span class="pre">float*</span> <span class="pre">buf_0</span> <span class="pre">=</span> <span class="pre">(float*)malloc(4</span> <span class="pre">*</span> <span class="pre">100);</span></code></p>
<p>As mentioned in the previous step, in addition to the subgraph input and output tensors, we may also need buffers to keep the intermediate results. To generate the buffer, we extract the shape information to determine the buffer type and size:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="c1">// This example only supports single output.</span>
<span class="k">auto</span> <span class="n">type_node</span> <span class="o">=</span> <span class="n">call</span><span class="o">-&gt;</span><span class="n">checked_type</span><span class="p">().</span><span class="n">as</span><span class="o">&lt;</span><span class="n">TensorTypeNode</span><span class="o">&gt;</span><span class="p">();</span>
<span class="n">ICHECK</span><span class="p">(</span><span class="n">type_node</span> <span class="o">!=</span> <span class="k">nullptr</span> <span class="o">&amp;&amp;</span> <span class="n">runtime</span><span class="o">::</span><span class="n">TypeMatch</span><span class="p">(</span><span class="n">type_node</span><span class="o">-&gt;</span><span class="n">dtype</span><span class="p">,</span> <span class="n">kDLFloat</span><span class="p">,</span> <span class="mi">32</span><span class="p">))</span>
      <span class="o">&lt;&lt;</span> <span class="s">&quot;Only support single output tensor with float type&quot;</span><span class="p">;</span>

<span class="c1">// Generate a unique buffer name.</span>
<span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">out</span> <span class="o">=</span> <span class="s">&quot;buf_&quot;</span> <span class="o">+</span> <span class="n">std</span><span class="o">::</span><span class="n">to_string</span><span class="p">(</span><span class="n">buf_idx_</span><span class="o">++</span><span class="p">);</span>

<span class="c1">// Extract the shape to be the buffer size.</span>
<span class="k">auto</span> <span class="n">out_shape</span> <span class="o">=</span> <span class="n">GetShape</span><span class="p">(</span><span class="n">call</span><span class="o">-&gt;</span><span class="n">checked_type</span><span class="p">());</span>
<span class="kt">int</span> <span class="n">out_size</span> <span class="o">=</span> <span class="mi">1</span><span class="p">;</span>
<span class="k">for</span> <span class="p">(</span><span class="kt">size_t</span> <span class="n">i</span> <span class="o">=</span> <span class="mi">0</span><span class="p">;</span> <span class="n">i</span> <span class="o">&lt;</span> <span class="n">out_shape</span><span class="p">.</span><span class="n">size</span><span class="p">();</span> <span class="o">++</span><span class="n">i</span><span class="p">)</span> <span class="p">{</span>
  <span class="n">out_size</span> <span class="o">*=</span> <span class="n">out_shape</span><span class="p">[</span><span class="n">i</span><span class="p">];</span>
<span class="p">}</span>

<span class="c1">// Make the buffer allocation and push to the buffer declarations.</span>
<span class="n">buf_stream</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;float* &quot;</span> <span class="o">&lt;&lt;</span> <span class="n">out</span> <span class="o">&lt;&lt;</span> <span class="s">&quot; = (float*)std::malloc(4 * &quot;</span> <span class="o">&lt;&lt;</span> <span class="n">out_size</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;);&quot;</span><span class="p">;</span>
<span class="n">buf_decl_</span><span class="p">.</span><span class="n">push_back</span><span class="p">(</span><span class="n">buf_stream</span><span class="p">.</span><span class="n">str</span><span class="p">());</span>
</pre></div>
</div>
<p>After we have allocated the output buffer, we can now close the function call string and push the generated function call to a class variable <code class="docutils literal notranslate"><span class="pre">ext_func_body</span></code>.</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">decl_stream</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;, &quot;</span> <span class="o">&lt;&lt;</span> <span class="n">out</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;);&quot;</span><span class="p">;</span>
<span class="n">ext_func_body</span><span class="p">.</span><span class="n">push_back</span><span class="p">(</span><span class="n">decl_stream</span><span class="p">.</span><span class="n">str</span><span class="p">());</span>
</pre></div>
</div>
<p><strong>4. Update output buffer</strong></p>
<p>To let the next node, which accepts the output of the current call node as its input, know which buffer it should take, we need to update the class variable <code class="docutils literal notranslate"><span class="pre">out_</span></code> before leaving this visit function:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">out_</span><span class="p">.</span><span class="n">clear</span><span class="p">();</span>
<span class="n">out_</span><span class="p">.</span><span class="n">push_back</span><span class="p">({</span><span class="n">out</span><span class="p">,</span> <span class="n">out_size</span><span class="p">});</span>
</pre></div>
</div>
<p>Congratulations! we have finished the most difficult function in this class. In the next two sections, we just need to make up some minor missing parts in this function.</p>
</div>
<div class="section" id="code-generation-for-input-variables">
<h4>Code Generation for Input Variables<a class="headerlink" href="#code-generation-for-input-variables" title="永久链接至标题">¶</a></h4>
<p>Recall that we collected the input buffer information by visiting the arguments of a call node (2nd step in the previous section), and handled the case when its argument is another call node (4th step). In this section, we demonstrate how to handle other nodes by taking <code class="docutils literal notranslate"><span class="pre">VarNode</span></code> as an example.</p>
<p><code class="docutils literal notranslate"><span class="pre">VarNode</span></code> represents input tensors in a model. The only but important information it has is a name hint (e.g., <code class="docutils literal notranslate"><span class="pre">data</span></code>, <code class="docutils literal notranslate"><span class="pre">weight</span></code>, etc). When visiting a <code class="docutils literal notranslate"><span class="pre">VarNode</span></code>, we simply update class variable <code class="docutils literal notranslate"><span class="pre">out_</span></code> to pass the name hint so that the descendant call nodes can generate the correct function call.</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="kt">void</span> <span class="nf">VisitExpr_</span><span class="p">(</span><span class="k">const</span> <span class="n">VarNode</span><span class="o">*</span> <span class="n">node</span><span class="p">)</span> <span class="p">{</span>
  <span class="n">ext_func_args_</span><span class="p">.</span><span class="n">push_back</span><span class="p">(</span><span class="n">node</span><span class="o">-&gt;</span><span class="n">name_hint</span><span class="p">());</span>
  <span class="n">out_</span><span class="p">.</span><span class="n">clear</span><span class="p">();</span>
  <span class="n">out_</span><span class="p">.</span><span class="n">push_back</span><span class="p">({</span><span class="n">node</span><span class="o">-&gt;</span><span class="n">name_hint</span><span class="p">(),</span> <span class="mi">0</span><span class="p">});</span>
<span class="p">}</span>
</pre></div>
</div>
<p>Note that in this example we assume the subgraph we are offloading has only call nodes and variable nodes. If your subgraphs contain other types of nodes, such as <code class="docutils literal notranslate"><span class="pre">TupleNode</span></code>, then you also need to visit them and bypass the output buffer information.</p>
</div>
<div class="section" id="code-emitting">
<h4>Code Emitting<a class="headerlink" href="#code-emitting" title="永久链接至标题">¶</a></h4>
<p>The final part in this codegen class is a <code class="docutils literal notranslate"><span class="pre">JIT</span></code> function that emits a C function for the subgraph and uses the C code we just generated as the function body. Remember, in addition to the subgraph function we generated in the previous sections, we also need a wrapper function with a unified argument for TVM runtime to invoke and pass data. Fortunately, the base class we inherited already provides an implementation, <code class="docutils literal notranslate"><span class="pre">JitImpl</span></code>, to generate the function. For example, we can invoke <code class="docutils literal notranslate"><span class="pre">JitImpl</span></code> as follows:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">JitImpl</span><span class="p">(</span><span class="s">&quot;gcc_0&quot;</span> <span class="cm">/* Subgraph symbol (ID) */</span><span class="p">,</span>
        <span class="p">{</span><span class="s">&quot;gcc_input0&quot;</span><span class="p">,</span> <span class="s">&quot;gcc_input1&quot;</span><span class="p">,</span> <span class="s">&quot;gcc_input2&quot;</span><span class="p">,</span> <span class="s">&quot;gcc_input3&quot;</span><span class="p">}</span> <span class="cm">/* Input arguments */</span><span class="p">,</span>
        <span class="p">{</span><span class="s">&quot;float *buf_0 = (float*)malloc(4 * 20)&quot;</span><span class="p">,</span> <span class="p">...}</span> <span class="cm">/* Buffer allocations */</span><span class="p">,</span>
        <span class="p">{</span><span class="s">&quot;gcc_0_2(gcc_input0, gcc_input1, buf_0);&quot;</span><span class="p">}</span> <span class="cm">/* Function body */</span><span class="p">,</span>
        <span class="p">{</span><span class="s">&quot;out&quot;</span><span class="p">}</span> <span class="cm">/* Output */</span><span class="p">);</span>
</pre></div>
</div>
<p>The above call will generate three functions (one from the TVM wrapper macro):</p>
<ol class="arabic simple">
<li><p>The subgraph function <code class="docutils literal notranslate"><span class="pre">gcc_0_</span></code> (with one more underline at the end of the function name) with all C code we generated to execute a subgraph.</p></li>
<li><p>The wrapper function <code class="docutils literal notranslate"><span class="pre">gcc_0__wrapper_</span></code> with a list of <code class="docutils literal notranslate"><span class="pre">DLTensor</span></code> arguments that casts data to the right type and invokes <code class="docutils literal notranslate"><span class="pre">gcc_0_</span></code>.</p></li>
<li><p>The TVM runtime compatible function <code class="docutils literal notranslate"><span class="pre">gcc_0</span></code> with TVM unified function arguments that unpacks TVM packed tensors and invokes <code class="docutils literal notranslate"><span class="pre">gcc_0__wrapper_</span></code>.</p></li>
</ol>
<p>Accordingly, the only thing we need in <code class="docutils literal notranslate"><span class="pre">JIT</span></code> implementation is passing all subgraph function code we generated to <code class="docutils literal notranslate"><span class="pre">JitImpl</span></code>:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">JIT</span><span class="p">()</span> <span class="p">{</span>
  <span class="c1">// Write function macros</span>
  <span class="k">for</span> <span class="p">(</span><span class="k">auto</span> <span class="nl">decl</span> <span class="p">:</span> <span class="n">func_decl_</span><span class="p">)</span> <span class="p">{</span>
    <span class="n">code_stream_</span> <span class="o">&lt;&lt;</span> <span class="n">decl</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
  <span class="p">}</span>
  <span class="k">return</span> <span class="n">JitImpl</span><span class="p">(</span><span class="n">ext_func_id_</span><span class="p">,</span> <span class="n">ext_func_args_</span><span class="p">,</span> <span class="n">buf_decl_</span><span class="p">,</span> <span class="n">ext_func_body</span><span class="p">,</span> <span class="n">out_</span><span class="p">);</span>
<span class="p">}</span>
</pre></div>
</div>
<p>All variables (<code class="docutils literal notranslate"><span class="pre">ext_func_id</span></code>, etc) we passed are class variables and were filled when we traversed the subgraph.</p>
</div>
</div>
<div class="section" id="implement-csourcecodegen">
<h3>Implement CSourceCodegen<a class="headerlink" href="#implement-csourcecodegen" title="永久链接至标题">¶</a></h3>
<p>Again, let’s create a class skeleton and implement the required functions. Note that it inherits <code class="docutils literal notranslate"><span class="pre">CSourceModuleCodegenBase</span></code></p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="k">class</span> <span class="nc">CSourceCodegen</span> <span class="o">:</span> <span class="k">public</span> <span class="n">CSourceModuleCodegenBase</span> <span class="p">{</span>
 <span class="k">public</span><span class="o">:</span>
  <span class="c1">// Pass a subgraph function, and generate the C code.</span>
  <span class="kt">void</span> <span class="n">GenCFunc</span><span class="p">(</span><span class="k">const</span> <span class="n">Function</span><span class="o">&amp;</span> <span class="n">func</span><span class="p">)</span> <span class="p">{</span> <span class="p">;</span> <span class="p">}</span>

  <span class="c1">// Use GenCFunc to generate the C code and wrap it as a C source module.</span>
  <span class="n">runtime</span><span class="o">::</span><span class="n">Module</span> <span class="n">CreateCSourceModule</span><span class="p">(</span><span class="k">const</span> <span class="n">NodeRef</span><span class="o">&amp;</span> <span class="n">ref</span><span class="p">)</span> <span class="k">override</span> <span class="p">{</span> <span class="p">;</span> <span class="p">}</span>

 <span class="k">private</span><span class="o">:</span>
  <span class="n">std</span><span class="o">::</span><span class="n">ostringstream</span> <span class="n">code_stream_</span><span class="p">;</span>
<span class="p">};</span>
</pre></div>
</div>
<div class="section" id="implement-gencfunc">
<h4>Implement GenCFunc<a class="headerlink" href="#implement-gencfunc" title="永久链接至标题">¶</a></h4>
<p><code class="docutils literal notranslate"><span class="pre">GenCFunc</span></code> simply uses the <code class="docutils literal notranslate"><span class="pre">CodegenC</span></code> we just implemented to traverse a Relay function (subgraph) and obtains the generated C code. The builtin function <code class="docutils literal notranslate"><span class="pre">GetExtSymbol</span></code> retrieves a unique symbol name (e.g., <code class="docutils literal notranslate"><span class="pre">gcc_0</span></code>) in the Relay function and we <strong>must</strong> use it as the C function name, because this symbol is going to be used for DSO runtime lookup.</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="kt">void</span> <span class="nf">GenCFunc</span><span class="p">(</span><span class="k">const</span> <span class="n">Function</span><span class="o">&amp;</span> <span class="n">func</span><span class="p">)</span> <span class="p">{</span>
  <span class="n">ICHECK</span><span class="p">(</span><span class="n">func</span><span class="p">.</span><span class="n">defined</span><span class="p">())</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;Input error: expect a Relay function.&quot;</span><span class="p">;</span>

  <span class="c1">// Record the external symbol for runtime lookup.</span>
  <span class="k">auto</span> <span class="n">sid</span> <span class="o">=</span> <span class="n">GetExtSymbol</span><span class="p">(</span><span class="n">func</span><span class="p">);</span>

  <span class="n">CodeGenC</span> <span class="n">builder</span><span class="p">(</span><span class="n">sid</span><span class="p">);</span>
  <span class="n">builder</span><span class="p">.</span><span class="n">VisitExpr</span><span class="p">(</span><span class="n">func</span><span class="o">-&gt;</span><span class="n">body</span><span class="p">);</span>
  <span class="n">code_stream_</span> <span class="o">&lt;&lt;</span> <span class="n">builder</span><span class="p">.</span><span class="n">JIT</span><span class="p">();</span>
<span class="p">}</span>
</pre></div>
</div>
</div>
<div class="section" id="implement-createcsourcemodule">
<h4>Implement CreateCSourceModule<a class="headerlink" href="#implement-createcsourcemodule" title="永久链接至标题">¶</a></h4>
<p>This function creates a runtime module for the external library. In this example, we create a CSourceModule that can be directly compiled and linked together with a TVM generated DSOModule. After you have implemented <code class="docutils literal notranslate"><span class="pre">CodegenC</span></code>, implementing this function is relatively straightforward:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">runtime</span><span class="o">::</span><span class="n">Module</span> <span class="n">CreateCSourceModule</span><span class="p">(</span><span class="k">const</span> <span class="n">NodeRef</span><span class="o">&amp;</span> <span class="n">ref</span><span class="p">)</span> <span class="k">override</span> <span class="p">{</span>
  <span class="c1">// Create headers</span>
  <span class="n">code_stream_</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;#include &lt;cstdint&gt;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
  <span class="n">code_stream_</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;#include &lt;iostream&gt;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
  <span class="n">code_stream_</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;#include &lt;cstdlib&gt;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
  <span class="n">code_stream_</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;#include &lt;stdio.h&gt;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
  <span class="n">code_stream_</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;#include &lt;cstring&gt;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
  <span class="n">code_stream_</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;#include &lt;tvm/runtime/c_runtime_api.h&gt;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
  <span class="n">code_stream_</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;#include &lt;dlpack/dlpack.h&gt;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>

  <span class="c1">// Append some common macro for operator definition.</span>
  <span class="k">const</span> <span class="kt">char</span><span class="o">*</span> <span class="n">operator_macro</span> <span class="o">=</span> <span class="sa">R</span><span class="s">&quot;</span><span class="dl">op_macro(</span><span class="s"></span>
<span class="s">  #define CSOURCE_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_)       \</span>
<span class="s">    extern &quot;C&quot; void p_ID_(float* a, float* b, float* out) { \</span>
<span class="s">      for (int64_t i = 0; i &lt; p_DIM1_; ++i) {               \</span>
<span class="s">        out[i] = a[i] p_OP_ b[i];                           \</span>
<span class="s">      }                                                     \</span>
<span class="s">    }</span>

<span class="s">  #define CSOURCE_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_)  \</span>
<span class="s">    extern &quot;C&quot; void p_ID_(float* a, float* b, float* out) {     \</span>
<span class="s">      for (int64_t i = 0; i &lt; p_DIM1_; ++i) {                   \</span>
<span class="s">        for (int64_t j = 0; j &lt; p_DIM2_; ++j) {                 \</span>
<span class="s">          int64_t k = i * p_DIM2_ + j;                          \</span>
<span class="s">          out[k] = a[k] p_OP_ b[k];                             \</span>
<span class="s">        }                                                       \</span>
<span class="s">      }                                                         \</span>
<span class="s">    }</span>
<span class="s">  </span><span class="dl">)op_macro</span><span class="s">&quot;</span><span class="p">;</span>

  <span class="n">code_stream_</span> <span class="o">&lt;&lt;</span> <span class="n">operator_macro</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;</span><span class="se">\n\n</span><span class="s">&quot;</span><span class="p">;</span>

  <span class="c1">// Generate C code for the subgraph.</span>
  <span class="k">if</span> <span class="p">(</span><span class="n">ref</span><span class="o">-&gt;</span><span class="n">IsInstance</span><span class="o">&lt;</span><span class="n">FunctionNode</span><span class="o">&gt;</span><span class="p">())</span> <span class="p">{</span>
    <span class="n">GenCFunc</span><span class="p">(</span><span class="n">Downcast</span><span class="o">&lt;</span><span class="n">Function</span><span class="o">&gt;</span><span class="p">(</span><span class="n">ref</span><span class="p">));</span>
  <span class="p">}</span> <span class="k">else</span> <span class="k">if</span> <span class="p">(</span><span class="n">ref</span><span class="o">-&gt;</span><span class="n">IsInstance</span><span class="o">&lt;</span><span class="n">relay</span><span class="o">::</span><span class="n">ModuleNode</span><span class="o">&gt;</span><span class="p">())</span> <span class="p">{</span>
    <span class="n">relay</span><span class="o">::</span><span class="n">Module</span> <span class="n">mod</span> <span class="o">=</span> <span class="n">Downcast</span><span class="o">&lt;</span><span class="n">relay</span><span class="o">::</span><span class="n">Module</span><span class="o">&gt;</span><span class="p">(</span><span class="n">ref</span><span class="p">);</span>
    <span class="k">for</span> <span class="p">(</span><span class="k">const</span> <span class="k">auto</span><span class="o">&amp;</span> <span class="nl">it</span> <span class="p">:</span> <span class="n">mod</span><span class="o">-&gt;</span><span class="n">functions</span><span class="p">)</span> <span class="p">{</span>
      <span class="n">GenCFunc</span><span class="p">(</span><span class="n">Downcast</span><span class="o">&lt;</span><span class="n">Function</span><span class="o">&gt;</span><span class="p">(</span><span class="n">it</span><span class="p">.</span><span class="n">second</span><span class="p">));</span>
    <span class="p">}</span>
  <span class="p">}</span> <span class="k">else</span> <span class="p">{</span>
    <span class="n">LOG</span><span class="p">(</span><span class="n">FATAL</span><span class="p">)</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;The input ref is expected to be a Relay function or module&quot;</span>
               <span class="o">&lt;&lt;</span> <span class="s">&quot;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
  <span class="p">}</span>

  <span class="c1">// Create a CSourceModule</span>
  <span class="k">const</span> <span class="k">auto</span><span class="o">*</span> <span class="n">pf</span> <span class="o">=</span> <span class="n">runtime</span><span class="o">::</span><span class="n">Registry</span><span class="o">::</span><span class="n">Get</span><span class="p">(</span><span class="s">&quot;module.csource_module_create&quot;</span><span class="p">);</span>
  <span class="n">ICHECK</span><span class="p">(</span><span class="n">pf</span> <span class="o">!=</span> <span class="k">nullptr</span><span class="p">)</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;Cannot find csource module to create the external runtime module&quot;</span><span class="p">;</span>
  <span class="k">return</span> <span class="p">(</span><span class="o">*</span><span class="n">pf</span><span class="p">)(</span><span class="n">code_stream_</span><span class="p">.</span><span class="n">str</span><span class="p">(),</span> <span class="s">&quot;cc&quot;</span><span class="p">);</span>
<span class="p">}</span>
</pre></div>
</div>
</div>
</div>
<div class="section" id="register-your-codegen">
<h3>Register Your Codegen<a class="headerlink" href="#register-your-codegen" title="永久链接至标题">¶</a></h3>
<p>The last step is registering your codegen to TVM backend. We first implement a simple function to invoke our codegen and generate a runtime module.</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">runtime</span><span class="o">::</span><span class="n">Module</span> <span class="n">CCompiler</span><span class="p">(</span><span class="k">const</span> <span class="n">NodeRef</span><span class="o">&amp;</span> <span class="n">ref</span><span class="p">)</span> <span class="p">{</span>
  <span class="n">CSourceCodegen</span> <span class="n">csource</span><span class="p">;</span>
  <span class="k">return</span> <span class="n">csource</span><span class="p">.</span><span class="n">CreateCSourceModule</span><span class="p">(</span><span class="n">ref</span><span class="p">);</span>
<span class="p">}</span>
</pre></div>
</div>
<p>Finally, we register this function to TVM backend:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">TVM_REGISTER_GLOBAL</span><span class="p">(</span><span class="s">&quot;relay.ext.ccompiler&quot;</span><span class="p">).</span><span class="n">set_body_typed</span><span class="p">(</span><span class="n">CCompiler</span><span class="p">);</span>
</pre></div>
</div>
<p>where <code class="docutils literal notranslate"><span class="pre">ccompiler</span></code> is a customized tag to let TVM know this is the codegen it should use to generate and offload subgraphs when the subgraph is annotated with <code class="docutils literal notranslate"><span class="pre">ccompiler</span></code>.</p>
<p>Finally, a good practice is to set up a CMake configuration flag to include your compiler only for your customers. We first create a cmake file: <code class="docutils literal notranslate"><span class="pre">cmake/modules/contrib/CODEGENC.cmake</span></code>:</p>
<div class="highlight-cmake notranslate"><div class="highlight"><pre><span></span><span class="nb">if</span><span class="p">(</span><span class="s">USE_CODEGENC</span><span class="p">)</span>
  <span class="nb">file</span><span class="p">(</span><span class="s">GLOB</span> <span class="s">CSOURCE_RELAY_CONTRIB_SRC</span> <span class="s">src/relay/backend/contrib/codegen_c/codegen.cc</span><span class="p">)</span>
  <span class="nb">list</span><span class="p">(</span><span class="s">APPEND</span> <span class="s">COMPILER_SRCS</span> <span class="o">${</span><span class="nv">CSOURCE_RELAY_CONTRIB_SRC</span><span class="o">}</span><span class="p">)</span>
<span class="nb">endif</span><span class="p">(</span><span class="s">USE_CODEGENC</span><span class="p">)</span>
</pre></div>
</div>
<p>So that users can configure whether to include your compiler when configuring TVM using <code class="docutils literal notranslate"><span class="pre">config.cmake</span></code>:</p>
<div class="highlight-cmake notranslate"><div class="highlight"><pre><span></span><span class="nb">set</span><span class="p">(</span><span class="s">USE_CODEGENC</span> <span class="s">ON</span><span class="p">)</span>
</pre></div>
</div>
</div>
</div>
<div class="section" id="implement-a-codegen-for-your-representation">
<h2>Implement a Codegen for Your Representation<a class="headerlink" href="#implement-a-codegen-for-your-representation" title="永久链接至标题">¶</a></h2>
<p>Although we have demonstrated how to implement a C codegen, your hardware may require other forms of graph representation, such as JSON. In this case, you could modify <code class="docutils literal notranslate"><span class="pre">CodegenC</span></code> class we have implemented to generate your own graph representation and implement a customized runtime module to let TVM runtime know how this graph representation should be executed.</p>
<p>To simplify, we define a graph representation named “ExampleJSON” in this guide. ExampleJSON does not mean the real JSON but just a simple representation for graphs without a control flow. For example, assuming we have the following subgraph named <code class="docutils literal notranslate"><span class="pre">subgraph_0</span></code>:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span> <span class="n">input0</span>
   <span class="o">|</span>
  <span class="n">add</span> <span class="o">&lt;--</span> <span class="n">input1</span>
   <span class="o">|</span>
<span class="n">subtract</span> <span class="o">&lt;--</span> <span class="n">input2</span>
   <span class="o">|</span>
<span class="n">multiply</span> <span class="o">&lt;--</span> <span class="n">input3</span>
   <span class="o">|</span>
  <span class="n">out</span>
</pre></div>
</div>
<p>Then the ExampleJON of this subgraph looks like:</p>
<div class="highlight-none notranslate"><div class="highlight"><pre><span></span>subgraph_0
  input 0 10 10
  input 1 10 10
  input 2 10 10
  input 3 10 10
  add 4 inputs: 0 1 shape: 10 10
  sub 5 inputs: 4 2 shape: 10 10
  mul 6 inputs: 5 3 shape: 10 10
</pre></div>
</div>
<p>The <code class="docutils literal notranslate"><span class="pre">input</span></code> keyword declares an input tensor with its ID and shape; while the other statements describes computations in <code class="docutils literal notranslate"><span class="pre">&lt;op&gt;</span> <span class="pre">&lt;output</span> <span class="pre">ID&gt;</span> <span class="pre">inputs:</span> <span class="pre">[input</span> <span class="pre">ID]</span> <span class="pre">shape:</span> <span class="pre">[shape]</span></code> syntax.</p>
<p>In this section, our goal is to implement the following customized TVM runtime module to execute ExampleJSON graphs.</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">runtime</span><span class="o">::</span><span class="n">Module</span> <span class="n">ExampleJsonCompiler</span><span class="p">(</span><span class="k">const</span> <span class="n">NodeRef</span><span class="o">&amp;</span> <span class="n">ref</span><span class="p">)</span> <span class="p">{</span>
    <span class="n">ExampleJsonCodeGen</span> <span class="n">codegen</span><span class="p">(</span><span class="n">ref</span><span class="p">);</span>
    <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">code</span> <span class="o">=</span> <span class="n">codegen</span><span class="p">.</span><span class="n">gen</span><span class="p">();</span> <span class="c1">// Note 1</span>
    <span class="k">const</span> <span class="k">auto</span><span class="o">*</span> <span class="n">pf</span> <span class="o">=</span> <span class="n">runtime</span><span class="o">::</span><span class="n">Registry</span><span class="o">::</span><span class="n">Get</span><span class="p">(</span><span class="s">&quot;module.examplejson_module_create&quot;</span><span class="p">);</span> <span class="c1">// Note 2</span>
    <span class="n">ICHECK</span><span class="p">(</span><span class="n">pf</span> <span class="o">!=</span> <span class="k">nullptr</span><span class="p">)</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;Cannot find ExampleJson module to create the external runtime module&quot;</span><span class="p">;</span>
    <span class="k">return</span> <span class="p">(</span><span class="o">*</span><span class="n">pf</span><span class="p">)(</span><span class="n">code</span><span class="p">);</span>
<span class="p">}</span>
<span class="n">TVM_REGISTER_GLOBAL</span><span class="p">(</span><span class="s">&quot;relay.ext.examplejsoncompiler&quot;</span><span class="p">).</span><span class="n">set_body_typed</span><span class="p">(</span><span class="n">ExampleJsonCompiler</span><span class="p">);</span>
</pre></div>
</div>
<p><strong>Note 1</strong>: We will implement a customized codegen later to generate a ExampleJSON code string by taking a subgraph.</p>
<p><strong>Note 2</strong>: This line obtains a pointer to a function for creating the customized runtime module. You can see that it takes subgraph code in ExampleJSON format we just generated and initializes a runtime module.</p>
<p>In the following sections, we are going to introduce 1) how to implement <code class="docutils literal notranslate"><span class="pre">ExampleJsonCodeGen</span></code> and 2) how to implement and register <code class="docutils literal notranslate"><span class="pre">examplejson_module_create</span></code>.</p>
<div class="section" id="implement-examplejsoncodegen">
<h3>Implement ExampleJsonCodeGen<a class="headerlink" href="#implement-examplejsoncodegen" title="永久链接至标题">¶</a></h3>
<p>Similar to the C codegen, we also derive <code class="docutils literal notranslate"><span class="pre">ExampleJsonCodeGen</span></code> from <code class="docutils literal notranslate"><span class="pre">ExprVisitor</span></code> to make use of visitor patterns for subgraph traversing. On the other hand, we do not have to inherit <code class="docutils literal notranslate"><span class="pre">CodegenCBase</span></code> because we do not need TVM C++ wrappers. The codegen class is implemented as follows:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="cp">#include</span> <span class="cpf">&lt;tvm/relay/expr_functor.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/relay/transform.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/relay/type.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/module.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/object.h&gt;</span><span class="cp"></span>

<span class="cp">#include</span> <span class="cpf">&lt;fstream&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;sstream&gt;</span><span class="cp"></span>

<span class="k">namespace</span> <span class="n">tvm</span> <span class="p">{</span>
<span class="k">namespace</span> <span class="n">relay</span> <span class="p">{</span>
<span class="k">namespace</span> <span class="n">contrib</span> <span class="p">{</span>

<span class="k">class</span> <span class="nc">ExampleJsonCodeGen</span> <span class="o">:</span> <span class="k">public</span> <span class="n">ExprVisitor</span> <span class="p">{</span>
  <span class="k">public</span><span class="o">:</span>
    <span class="k">explicit</span> <span class="n">ExampleJsonCodeGen</span><span class="p">();</span>

    <span class="c1">// Note 1</span>
    <span class="kt">void</span> <span class="nf">VisitExpr_</span><span class="p">(</span><span class="k">const</span> <span class="n">VarNode</span><span class="o">*</span> <span class="n">node</span><span class="p">)</span> <span class="p">{</span> <span class="cm">/* Skip in this example. */</span> <span class="p">}</span>
    <span class="kt">void</span> <span class="nf">VisitExpr_</span><span class="p">(</span><span class="k">const</span> <span class="n">CallNode</span><span class="o">*</span> <span class="n">call</span><span class="p">)</span> <span class="k">final</span> <span class="p">{</span> <span class="cm">/* Skip in this example. */</span> <span class="p">}</span>

    <span class="c1">// Note 2</span>
    <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">gen</span><span class="p">(</span><span class="n">NodeRef</span><span class="o">&amp;</span> <span class="n">ref</span><span class="p">)</span> <span class="p">{</span>
        <span class="k">this</span><span class="o">-&gt;</span><span class="n">code</span> <span class="o">=</span> <span class="s">&quot;&quot;</span><span class="p">;</span>
        <span class="k">if</span> <span class="p">(</span><span class="n">ref</span><span class="o">-&gt;</span><span class="n">IsInstance</span><span class="o">&lt;</span><span class="n">FunctionNode</span><span class="o">&gt;</span><span class="p">())</span> <span class="p">{</span>
            <span class="k">this</span><span class="o">-&gt;</span><span class="n">visit</span><span class="p">(</span><span class="n">Downcast</span><span class="o">&lt;</span><span class="n">Function</span><span class="o">&gt;</span><span class="p">(</span><span class="n">ref</span><span class="p">));</span>
        <span class="p">}</span> <span class="k">else</span> <span class="k">if</span> <span class="p">(</span><span class="n">ref</span><span class="o">-&gt;</span><span class="n">IsInstance</span><span class="o">&lt;</span><span class="n">relay</span><span class="o">::</span><span class="n">ModuleNode</span><span class="o">&gt;</span><span class="p">())</span> <span class="p">{</span>
            <span class="n">relay</span><span class="o">::</span><span class="n">Module</span> <span class="n">mod</span> <span class="o">=</span> <span class="n">Downcast</span><span class="o">&lt;</span><span class="n">relay</span><span class="o">::</span><span class="n">Module</span><span class="o">&gt;</span><span class="p">(</span><span class="n">ref</span><span class="p">);</span>
            <span class="k">for</span> <span class="p">(</span><span class="k">const</span> <span class="k">auto</span><span class="o">&amp;</span> <span class="nl">it</span> <span class="p">:</span> <span class="n">mod</span><span class="o">-&gt;</span><span class="n">functions</span><span class="p">)</span> <span class="p">{</span>
                <span class="k">this</span><span class="o">-&gt;</span><span class="n">visit</span><span class="p">(</span><span class="n">Downcast</span><span class="o">&lt;</span><span class="n">Function</span><span class="o">&gt;</span><span class="p">(</span><span class="n">it</span><span class="p">.</span><span class="n">second</span><span class="p">));</span>
            <span class="p">}</span>
        <span class="p">}</span> <span class="k">else</span> <span class="p">{</span>
            <span class="n">LOG</span><span class="p">(</span><span class="n">FATAL</span><span class="p">)</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;The input ref is expected to be a Relay function or module&quot;</span><span class="p">;</span>
        <span class="p">}</span>
        <span class="k">return</span> <span class="k">this</span><span class="o">-&gt;</span><span class="n">code</span><span class="p">;</span>
    <span class="p">}</span>

  <span class="k">private</span><span class="o">:</span>
      <span class="cm">/*! \brief The function id that represents a C source function. */</span>
     <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">code</span><span class="p">;</span>
<span class="p">}</span>
</pre></div>
</div>
<p><strong>Note 1</strong>: We again implement corresponding visitor functions to generate ExampleJSON code and store it to a class variable <code class="docutils literal notranslate"><span class="pre">code</span></code> (we skip the visitor function implementation in this example as their concepts are basically the same as C codegen). After finished the graph visiting, we should have an ExampleJSON graph in <code class="docutils literal notranslate"><span class="pre">code</span></code>.</p>
<p><strong>Note 2</strong>: We define an internal API <code class="docutils literal notranslate"><span class="pre">gen</span></code> to take a subgraph and generate a ExampleJSON code. This API can be in an arbitrary name you prefer.</p>
<p>The next step is to implement a customized runtime to make use of the output of <code class="docutils literal notranslate"><span class="pre">ExampleJsonCodeGen</span></code>.</p>
</div>
<div class="section" id="implement-a-customized-runtime">
<h3>Implement a Customized Runtime<a class="headerlink" href="#implement-a-customized-runtime" title="永久链接至标题">¶</a></h3>
<p>In this section, we will implement a customized TVM runtime step-by-step and register it to TVM runtime modules. The customized runtime should be located at <code class="docutils literal notranslate"><span class="pre">src/runtime/contrib/&lt;your-runtime-name&gt;/</span></code>. In our example, we name our runtime “example_ext_runtime”.</p>
<p>Again, we first define a customized runtime class as follows. The class has to be derived from TVM <code class="docutils literal notranslate"><span class="pre">ModuleNode</span></code> in order to be compatible with other TVM runtime modules.</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="cp">#include</span> <span class="cpf">&lt;dmlc/logging.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/c_runtime_api.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/memory.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/module.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/ndarray.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/object.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/packed_func.h&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;tvm/runtime/registry.h&gt;</span><span class="cp"></span>

<span class="cp">#include</span> <span class="cpf">&lt;fstream&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;cmath&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;map&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;sstream&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;string&gt;</span><span class="cp"></span>
<span class="cp">#include</span> <span class="cpf">&lt;vector&gt;</span><span class="cp"></span>

<span class="k">namespace</span> <span class="n">tvm</span> <span class="p">{</span>
<span class="k">namespace</span> <span class="n">runtime</span> <span class="p">{</span>
<span class="k">class</span> <span class="nc">ExampleJsonModule</span> <span class="o">:</span> <span class="k">public</span> <span class="n">ModuleNode</span> <span class="p">{</span>
 <span class="k">public</span><span class="o">:</span>
  <span class="k">explicit</span> <span class="n">ExampleJsonModule</span><span class="p">(</span><span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">graph_json</span><span class="p">);</span>

  <span class="n">PackedFunc</span> <span class="nf">GetFunction</span><span class="p">(</span><span class="k">const</span> <span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&amp;</span> <span class="n">name</span><span class="p">,</span>
                         <span class="k">const</span> <span class="n">ObjectPtr</span><span class="o">&lt;</span><span class="n">Object</span><span class="o">&gt;&amp;</span> <span class="n">sptr_to_self</span><span class="p">)</span> <span class="k">final</span><span class="p">;</span>

  <span class="k">const</span> <span class="kt">char</span><span class="o">*</span> <span class="nf">type_key</span><span class="p">()</span> <span class="k">const</span> <span class="p">{</span> <span class="k">return</span> <span class="s">&quot;examplejson&quot;</span><span class="p">;</span> <span class="p">}</span>

  <span class="kt">void</span> <span class="nf">SaveToBinary</span><span class="p">(</span><span class="n">dmlc</span><span class="o">::</span><span class="n">Stream</span><span class="o">*</span> <span class="n">stream</span><span class="p">)</span> <span class="k">final</span><span class="p">;</span>

  <span class="k">static</span> <span class="n">Module</span> <span class="nf">LoadFromBinary</span><span class="p">(</span><span class="kt">void</span><span class="o">*</span> <span class="n">strm</span><span class="p">);</span>

  <span class="k">static</span> <span class="n">Module</span> <span class="nf">Create</span><span class="p">(</span><span class="k">const</span> <span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&amp;</span> <span class="n">path</span><span class="p">);</span>

  <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">GetSource</span><span class="p">(</span><span class="k">const</span> <span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&amp;</span> <span class="n">format</span> <span class="o">=</span> <span class="s">&quot;&quot;</span><span class="p">);</span>

  <span class="kt">void</span> <span class="nf">Run</span><span class="p">(</span><span class="kt">int</span> <span class="n">id</span><span class="p">,</span> <span class="k">const</span> <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="kt">int</span><span class="o">&gt;&amp;</span> <span class="n">inputs</span><span class="p">,</span> <span class="kt">int</span> <span class="n">output</span><span class="p">);</span>

  <span class="kt">void</span> <span class="nf">ParseJson</span><span class="p">(</span><span class="k">const</span> <span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&amp;</span> <span class="n">json</span><span class="p">);</span>

 <span class="k">private</span><span class="o">:</span>
  <span class="cm">/* \brief The json string that represents a computational graph. */</span>
  <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">graph_json_</span><span class="p">;</span>
  <span class="cm">/* \brief The subgraph that being processed. */</span>
  <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">curr_subgraph_</span><span class="p">;</span>
  <span class="cm">/*! \brief A simple graph from subgraph id to node entries. */</span>
  <span class="n">std</span><span class="o">::</span><span class="n">map</span><span class="o">&lt;</span><span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="p">,</span> <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="n">NodeEntry</span><span class="o">&gt;</span> <span class="o">&gt;</span> <span class="n">graph_</span><span class="p">;</span>
  <span class="cm">/* \brief A simple pool to contain the tensor for each node in the graph. */</span>
  <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="n">NDArray</span><span class="o">&gt;</span> <span class="n">data_entry_</span><span class="p">;</span>
  <span class="cm">/* \brief A mapping from node id to op name. */</span>
  <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&gt;</span> <span class="n">op_id_</span><span class="p">;</span>
<span class="p">};</span>
</pre></div>
</div>
<p>In particular, there are some functions derived from <code class="docutils literal notranslate"><span class="pre">ModuleNode</span></code> that we must implement in <code class="docutils literal notranslate"><span class="pre">ExampleJsonModule</span></code>:</p>
<ul class="simple">
<li><p>Constructor: The constructor of this class should accept a subgraph (in your representation), process and store it in any format you like. The saved subgraph could be used by the following two functions.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">GetFunction</span></code>: This is the most important function in this class. When TVM runtime wants to execute a subgraph with your compiler tag, TVM runtime invokes this function from your customized runtime module. It provides the function name as well as runtime arguments, and <code class="docutils literal notranslate"><span class="pre">GetFunction</span></code> should return a packed function implementation for TVM runtime to execute.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">SaveToBinary</span></code> and <code class="docutils literal notranslate"><span class="pre">LoadFromBinary</span></code>: <code class="docutils literal notranslate"><span class="pre">SaveToBinary</span></code> serialize the runtime module to a binary format for later deployment. This function will be called by TVM when users use <code class="docutils literal notranslate"><span class="pre">export_library</span></code> API. On the other hand, since we are now using our own graph representation, we have to make sure that <code class="docutils literal notranslate"><span class="pre">LoadFromBinary</span></code> is able to construct the same runtime module by taking the serialized binary generated by <code class="docutils literal notranslate"><span class="pre">SaveToBinary</span></code>.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">GetSource</span></code> (optional): If you would like to see the generated ExampleJSON code, you can implement this function to dump it; otherwise you can skip the implementation.</p></li>
</ul>
<p>Other functions and class variables will be introduced along with the implementation of above must-have functions.</p>
<div class="section" id="implement-constructor">
<h4>Implement Constructor<a class="headerlink" href="#implement-constructor" title="永久链接至标题">¶</a></h4>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="k">explicit</span> <span class="nf">ExampleJsonModule</span><span class="p">(</span><span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">graph_json</span><span class="p">)</span> <span class="p">{</span>
  <span class="k">this</span><span class="o">-&gt;</span><span class="n">graph_json_</span> <span class="o">=</span> <span class="n">graph_json</span><span class="p">;</span>
  <span class="n">ParseJson</span><span class="p">(</span><span class="k">this</span><span class="o">-&gt;</span><span class="n">graph_json_</span><span class="p">);</span>
<span class="p">}</span>
</pre></div>
</div>
<p>Then, we implement <code class="docutils literal notranslate"><span class="pre">ParseJson</span></code> to parse a subgraph in ExampleJSON format and construct a graph in memory for later usage. Since we do not support subgraph with branches in this example, we simply use an array to store every nodes in a subgraph in order.</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="kt">void</span> <span class="nf">ParseJson</span><span class="p">(</span><span class="k">const</span> <span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&amp;</span> <span class="n">json</span><span class="p">)</span> <span class="p">{</span>
  <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">line</span><span class="p">;</span>
  <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">curr_subgraph</span><span class="p">;</span>
  <span class="n">std</span><span class="o">::</span><span class="n">stringstream</span> <span class="n">ss</span><span class="p">(</span><span class="n">json</span><span class="p">);</span>

  <span class="k">while</span> <span class="p">(</span><span class="n">std</span><span class="o">::</span><span class="n">getline</span><span class="p">(</span><span class="n">ss</span><span class="p">,</span> <span class="n">line</span><span class="p">,</span> <span class="sc">&#39;\n&#39;</span><span class="p">))</span> <span class="p">{</span>
    <span class="n">std</span><span class="o">::</span><span class="n">stringstream</span> <span class="n">ss2</span><span class="p">(</span><span class="n">line</span><span class="p">);</span>
    <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">token</span><span class="p">;</span>
    <span class="kt">int</span> <span class="n">id</span> <span class="o">=</span> <span class="mi">0</span><span class="p">;</span>

    <span class="n">ss2</span> <span class="o">&gt;&gt;</span> <span class="n">token</span><span class="p">;</span>
    <span class="k">if</span> <span class="p">(</span><span class="n">token</span><span class="p">.</span><span class="n">find</span><span class="p">(</span><span class="s">&quot;subgraph_&quot;</span><span class="p">)</span> <span class="o">!=</span> <span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">::</span><span class="n">npos</span><span class="p">)</span> <span class="p">{</span>
      <span class="n">curr_subgraph</span> <span class="o">=</span> <span class="n">token</span><span class="p">;</span>
      <span class="k">continue</span><span class="p">;</span>
    <span class="p">}</span>

    <span class="n">ss2</span> <span class="o">&gt;&gt;</span> <span class="n">id</span><span class="p">;</span>
    <span class="k">if</span> <span class="p">(</span><span class="n">op_id_</span><span class="p">.</span><span class="n">size</span><span class="p">()</span> <span class="o">&lt;=</span> <span class="k">static_cast</span><span class="o">&lt;</span><span class="kt">size_t</span><span class="o">&gt;</span><span class="p">(</span><span class="n">id</span><span class="p">))</span> <span class="p">{</span>
      <span class="n">op_id_</span><span class="p">.</span><span class="n">resize</span><span class="p">(</span><span class="n">id</span> <span class="o">+</span> <span class="mi">1</span><span class="p">);</span>
      <span class="n">data_entry_</span><span class="p">.</span><span class="n">resize</span><span class="p">(</span><span class="n">id</span> <span class="o">+</span> <span class="mi">1</span><span class="p">);</span>
    <span class="p">}</span>

    <span class="kt">int64_t</span> <span class="n">total_elements</span> <span class="o">=</span> <span class="mi">1</span><span class="p">;</span>
    <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="kt">int64_t</span><span class="o">&gt;</span> <span class="n">shape</span><span class="p">;</span>
    <span class="k">if</span> <span class="p">(</span><span class="n">token</span> <span class="o">==</span> <span class="s">&quot;input&quot;</span><span class="p">)</span> <span class="p">{</span>
      <span class="kt">int64_t</span> <span class="n">size</span> <span class="o">=</span> <span class="mi">0</span><span class="p">;</span>
      <span class="k">while</span> <span class="p">(</span><span class="n">ss2</span> <span class="o">&gt;&gt;</span> <span class="n">size</span><span class="p">)</span> <span class="p">{</span>
        <span class="n">total_elements</span> <span class="o">*=</span> <span class="n">size</span><span class="p">;</span>
        <span class="n">shape</span><span class="p">.</span><span class="n">push_back</span><span class="p">(</span><span class="n">size</span><span class="p">);</span>
      <span class="p">}</span>
    <span class="p">}</span> <span class="k">else</span> <span class="p">{</span>
      <span class="n">op_id_</span><span class="p">[</span><span class="n">id</span><span class="p">]</span> <span class="o">=</span> <span class="n">token</span><span class="p">;</span> <span class="c1">// Note 1</span>
      <span class="kt">bool</span> <span class="n">shape_data</span> <span class="o">=</span> <span class="nb">false</span><span class="p">;</span>
      <span class="n">NodeEntry</span> <span class="n">entry</span><span class="p">;</span>
      <span class="k">while</span> <span class="p">(</span><span class="n">ss2</span> <span class="o">&gt;&gt;</span> <span class="n">token</span><span class="p">)</span> <span class="p">{</span>
        <span class="k">if</span> <span class="p">(</span><span class="n">token</span> <span class="o">==</span> <span class="s">&quot;shape:&quot;</span><span class="p">)</span> <span class="p">{</span>
          <span class="n">shape_data</span> <span class="o">=</span> <span class="nb">true</span><span class="p">;</span>
        <span class="p">}</span> <span class="k">else</span> <span class="k">if</span> <span class="p">(</span><span class="n">shape_data</span><span class="p">)</span> <span class="p">{</span>
          <span class="n">total_elements</span> <span class="o">*=</span> <span class="n">std</span><span class="o">::</span><span class="n">stoll</span><span class="p">(</span><span class="n">token</span><span class="p">);</span>
          <span class="n">shape</span><span class="p">.</span><span class="n">push_back</span><span class="p">(</span><span class="n">std</span><span class="o">::</span><span class="n">stoll</span><span class="p">(</span><span class="n">token</span><span class="p">));</span>
        <span class="p">}</span> <span class="k">else</span> <span class="k">if</span> <span class="p">(</span><span class="n">token</span> <span class="o">!=</span> <span class="s">&quot;inputs:&quot;</span><span class="p">)</span> <span class="p">{</span>
          <span class="n">entry</span><span class="p">.</span><span class="n">inputs</span><span class="p">.</span><span class="n">push_back</span><span class="p">(</span><span class="n">std</span><span class="o">::</span><span class="n">stoi</span><span class="p">(</span><span class="n">token</span><span class="p">));</span>
        <span class="p">}</span>
      <span class="p">}</span>
      <span class="n">entry</span><span class="p">.</span><span class="n">id</span> <span class="o">=</span> <span class="n">id</span><span class="p">;</span>
      <span class="n">entry</span><span class="p">.</span><span class="n">output</span> <span class="o">=</span> <span class="n">id</span><span class="p">;</span>
      <span class="n">graph_</span><span class="p">[</span><span class="n">curr_subgraph</span><span class="p">].</span><span class="n">push_back</span><span class="p">(</span><span class="n">entry</span><span class="p">);</span> <span class="c1">// Note 2</span>
    <span class="p">}</span>
    <span class="n">DLDevice</span> <span class="n">dev</span><span class="p">;</span>
    <span class="n">dev</span><span class="p">.</span><span class="n">device_type</span> <span class="o">=</span> <span class="k">static_cast</span><span class="o">&lt;</span><span class="n">DLDeviceType</span><span class="o">&gt;</span><span class="p">(</span><span class="mi">1</span><span class="p">);</span>
    <span class="n">dev</span><span class="p">.</span><span class="n">device_id</span> <span class="o">=</span> <span class="mi">0</span><span class="p">;</span>
    <span class="n">data_entry_</span><span class="p">[</span><span class="n">id</span><span class="p">]</span> <span class="o">=</span> <span class="n">NDArray</span><span class="o">::</span><span class="n">Empty</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="n">DLDataType</span><span class="p">{</span><span class="n">kDLFloat</span><span class="p">,</span> <span class="mi">32</span><span class="p">,</span> <span class="mi">1</span><span class="p">},</span> <span class="n">dev</span><span class="p">);</span> <span class="c1">// Note 3</span>
  <span class="p">}</span>
<span class="p">}</span>
</pre></div>
</div>
<p><strong>Note 1</strong>: We use a class variable <code class="docutils literal notranslate"><span class="pre">op_id_</span></code> to map from subgraph node ID to the operator name (e.g., <code class="docutils literal notranslate"><span class="pre">add</span></code>) so that we can invoke the corresponding operator function in runtime.</p>
<p><strong>Note 2</strong>: We use a class variable <code class="docutils literal notranslate"><span class="pre">graph_</span></code> to map from subgraph name to an array of nodes. <code class="docutils literal notranslate"><span class="pre">GetFunction</span></code> will query graph nodes by a subgraph ID in runtime.</p>
<p><strong>Note 3</strong>: We use a class variable <cite>data_entry_</cite> to map from a subgraph node ID to a tensor data placeholder. We will put inputs and outputs to the corresponding data entry in runtime.</p>
</div>
<div class="section" id="implement-getfunction">
<h4>Implement GetFunction<a class="headerlink" href="#implement-getfunction" title="永久链接至标题">¶</a></h4>
<p>After the construction, we should have the above class variables ready. We then implement <code class="docutils literal notranslate"><span class="pre">GetFunction</span></code> to provide executable subgraph functions to TVM runtime:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">PackedFunc</span> <span class="nf">GetFunction</span><span class="p">(</span><span class="k">const</span> <span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&amp;</span> <span class="n">name</span><span class="p">,</span>
                       <span class="k">const</span> <span class="n">ObjectPtr</span><span class="o">&lt;</span><span class="n">Object</span><span class="o">&gt;&amp;</span> <span class="n">sptr_to_self</span><span class="p">)</span> <span class="k">final</span> <span class="p">{</span>
  <span class="k">if</span> <span class="p">(</span><span class="k">this</span><span class="o">-&gt;</span><span class="n">graph_</span><span class="p">.</span><span class="n">find</span><span class="p">(</span><span class="n">name</span><span class="p">)</span> <span class="o">!=</span> <span class="k">this</span><span class="o">-&gt;</span><span class="n">graph_</span><span class="p">.</span><span class="n">end</span><span class="p">())</span> <span class="p">{</span>
    <span class="k">this</span><span class="o">-&gt;</span><span class="n">curr_subgraph_</span> <span class="o">=</span> <span class="n">name</span><span class="p">;</span>
    <span class="k">return</span> <span class="n">PackedFunc</span><span class="p">([</span><span class="n">sptr_to_self</span><span class="p">,</span> <span class="k">this</span><span class="p">](</span><span class="n">TVMArgs</span> <span class="n">args</span><span class="p">,</span> <span class="n">TVMRetValue</span><span class="o">*</span> <span class="n">rv</span><span class="p">)</span> <span class="p">{</span>

      <span class="c1">// Copy input tensors to corresponding data entries.</span>
      <span class="k">for</span> <span class="p">(</span><span class="k">auto</span> <span class="n">i</span> <span class="o">=</span> <span class="mi">0</span><span class="p">;</span> <span class="n">i</span> <span class="o">&lt;</span> <span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">();</span> <span class="o">++</span><span class="n">i</span><span class="p">)</span> <span class="p">{</span>
        <span class="n">ICHECK</span><span class="p">(</span><span class="n">args</span><span class="p">[</span><span class="n">i</span><span class="p">].</span><span class="n">type_code</span><span class="p">()</span> <span class="o">==</span> <span class="n">kNDArrayContainer</span> <span class="o">||</span> <span class="n">args</span><span class="p">[</span><span class="n">i</span><span class="p">].</span><span class="n">type_code</span><span class="p">()</span> <span class="o">==</span> <span class="n">kArrayHandle</span><span class="p">)</span>
            <span class="o">&lt;&lt;</span> <span class="s">&quot;Expect NDArray or DLTensor as inputs</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
        <span class="k">if</span> <span class="p">(</span><span class="n">args</span><span class="p">[</span><span class="n">i</span><span class="p">].</span><span class="n">type_code</span><span class="p">()</span> <span class="o">==</span> <span class="n">kArrayHandle</span><span class="p">)</span> <span class="p">{</span>
          <span class="n">DLTensor</span><span class="o">*</span> <span class="n">arg</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="n">i</span><span class="p">];</span>
          <span class="k">this</span><span class="o">-&gt;</span><span class="n">data_entry_</span><span class="p">[</span><span class="n">i</span><span class="p">].</span><span class="n">CopyFrom</span><span class="p">(</span><span class="n">arg</span><span class="p">);</span>
        <span class="p">}</span> <span class="k">else</span> <span class="p">{</span>
          <span class="n">NDArray</span> <span class="n">arg</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="n">i</span><span class="p">];</span>
          <span class="k">this</span><span class="o">-&gt;</span><span class="n">data_entry_</span><span class="p">[</span><span class="n">i</span><span class="p">].</span><span class="n">CopyFrom</span><span class="p">(</span><span class="n">arg</span><span class="p">);</span>
        <span class="p">}</span>
      <span class="p">}</span>

      <span class="c1">// Execute the subgraph.</span>
      <span class="k">for</span> <span class="p">(</span><span class="k">const</span> <span class="k">auto</span><span class="o">&amp;</span> <span class="nl">it</span> <span class="p">:</span> <span class="k">this</span><span class="o">-&gt;</span><span class="n">graph_</span><span class="p">[</span><span class="k">this</span><span class="o">-&gt;</span><span class="n">curr_subgraph_</span><span class="p">])</span> <span class="p">{</span>
        <span class="k">this</span><span class="o">-&gt;</span><span class="n">Run</span><span class="p">(</span><span class="n">it</span><span class="p">.</span><span class="n">id</span><span class="p">,</span> <span class="n">it</span><span class="p">.</span><span class="n">inputs</span><span class="p">,</span> <span class="n">it</span><span class="p">.</span><span class="n">output</span><span class="p">);</span>
      <span class="p">}</span>
      <span class="n">ICHECK_GT</span><span class="p">(</span><span class="n">graph_</span><span class="p">.</span><span class="n">count</span><span class="p">(</span><span class="k">this</span><span class="o">-&gt;</span><span class="n">curr_subgraph_</span><span class="p">),</span> <span class="mi">0U</span><span class="p">);</span>

      <span class="c1">// Copy the output from a data entry back to TVM runtime argument.</span>
      <span class="k">auto</span> <span class="n">out_idx</span> <span class="o">=</span> <span class="n">graph_</span><span class="p">[</span><span class="k">this</span><span class="o">-&gt;</span><span class="n">curr_subgraph_</span><span class="p">].</span><span class="n">back</span><span class="p">().</span><span class="n">output</span><span class="p">;</span>
      <span class="k">if</span> <span class="p">(</span><span class="n">args</span><span class="p">[</span><span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">()</span> <span class="o">-</span> <span class="mi">1</span><span class="p">].</span><span class="n">type_code</span><span class="p">()</span> <span class="o">==</span> <span class="n">kArrayHandle</span><span class="p">)</span> <span class="p">{</span>
        <span class="n">DLTensor</span><span class="o">*</span> <span class="n">arg</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">()</span> <span class="o">-</span> <span class="mi">1</span><span class="p">];</span>
        <span class="k">this</span><span class="o">-&gt;</span><span class="n">data_entry_</span><span class="p">[</span><span class="n">out_idx</span><span class="p">].</span><span class="n">CopyTo</span><span class="p">(</span><span class="n">arg</span><span class="p">);</span>
      <span class="p">}</span> <span class="k">else</span> <span class="p">{</span>
        <span class="n">NDArray</span> <span class="n">arg</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">()</span> <span class="o">-</span> <span class="mi">1</span><span class="p">];</span>
        <span class="k">this</span><span class="o">-&gt;</span><span class="n">data_entry_</span><span class="p">[</span><span class="n">out_idx</span><span class="p">].</span><span class="n">CopyTo</span><span class="p">(</span><span class="n">arg</span><span class="p">);</span>
      <span class="p">}</span>
      <span class="o">*</span><span class="n">rv</span> <span class="o">=</span> <span class="n">data_entry_</span><span class="p">.</span><span class="n">back</span><span class="p">();</span>
    <span class="p">});</span>
  <span class="p">}</span> <span class="k">else</span> <span class="p">{</span>
    <span class="n">LOG</span><span class="p">(</span><span class="n">FATAL</span><span class="p">)</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;Unknown subgraph: &quot;</span> <span class="o">&lt;&lt;</span> <span class="n">name</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
    <span class="k">return</span> <span class="n">PackedFunc</span><span class="p">();</span>
  <span class="p">}</span>
<span class="p">}</span>
</pre></div>
</div>
<p>As can be seen, <code class="docutils literal notranslate"><span class="pre">GetFunction</span></code> is composed of three major parts. The first part copies data from TVM runtime arguments to the corresponding data entries we assigned in the constructor. The second part executes the subgraph with <code class="docutils literal notranslate"><span class="pre">Run</span></code> function (will implement later) and saves the results to another data entry. The third part copies the results from the output data entry back to the corresponding TVM runtime argument for output.</p>
</div>
<div class="section" id="implement-run">
<h4>Implement Run<a class="headerlink" href="#implement-run" title="永久链接至标题">¶</a></h4>
<p>Now let’s implement <code class="docutils literal notranslate"><span class="pre">Run</span></code> function. This function accepts 1) a subgraph ID, 2) a list of input data entry indexs, and 3) an output data entry index.</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="kt">void</span> <span class="nf">Run</span><span class="p">(</span><span class="kt">int</span> <span class="n">id</span><span class="p">,</span> <span class="k">const</span> <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="kt">int</span><span class="o">&gt;&amp;</span> <span class="n">inputs</span><span class="p">,</span> <span class="kt">int</span> <span class="n">output</span><span class="p">)</span> <span class="p">{</span>
  <span class="c1">// Make a list data entry indexs.</span>
  <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="kt">int</span><span class="o">&gt;</span> <span class="n">args</span><span class="p">(</span><span class="n">inputs</span><span class="p">.</span><span class="n">begin</span><span class="p">(),</span> <span class="n">inputs</span><span class="p">.</span><span class="n">end</span><span class="p">());</span>
  <span class="n">args</span><span class="p">.</span><span class="n">push_back</span><span class="p">(</span><span class="n">output</span><span class="p">);</span>

  <span class="c1">// Initialize data holders.</span>
  <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="n">TVMValue</span><span class="o">&gt;</span> <span class="n">values</span><span class="p">(</span><span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">());</span>
  <span class="n">std</span><span class="o">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="kt">int</span><span class="o">&gt;</span> <span class="n">type_codes</span><span class="p">(</span><span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">());</span>

  <span class="c1">// Initialize a TVM arg setter with TVMValue and its type code.</span>
  <span class="n">TVMArgsSetter</span> <span class="n">setter</span><span class="p">(</span><span class="n">values</span><span class="p">.</span><span class="n">data</span><span class="p">(),</span> <span class="n">type_codes</span><span class="p">.</span><span class="n">data</span><span class="p">());</span>

  <span class="c1">// Set each argument to its corresponding data entry.</span>
  <span class="k">if</span> <span class="p">(</span><span class="n">op_id_</span><span class="p">[</span><span class="n">id</span><span class="p">]</span> <span class="o">==</span> <span class="s">&quot;add&quot;</span> <span class="o">||</span> <span class="n">op_id_</span><span class="p">[</span><span class="n">id</span><span class="p">]</span> <span class="o">==</span> <span class="s">&quot;sub&quot;</span> <span class="o">||</span> <span class="n">op_id_</span><span class="p">[</span><span class="n">id</span><span class="p">]</span> <span class="o">==</span> <span class="s">&quot;mul&quot;</span><span class="p">)</span> <span class="p">{</span>
    <span class="k">for</span> <span class="p">(</span><span class="kt">size_t</span> <span class="n">i</span> <span class="o">=</span> <span class="mi">0</span><span class="p">;</span> <span class="n">i</span> <span class="o">&lt;</span> <span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">();</span> <span class="n">i</span><span class="o">++</span><span class="p">)</span> <span class="p">{</span>
      <span class="n">setter</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">data_entry_</span><span class="p">[</span><span class="n">args</span><span class="p">[</span><span class="n">i</span><span class="p">]]);</span>
    <span class="p">}</span>
  <span class="p">}</span>

  <span class="c1">// Invoke the corresponding operator function.</span>
  <span class="k">if</span> <span class="p">(</span><span class="n">op_id_</span><span class="p">[</span><span class="n">id</span><span class="p">]</span> <span class="o">==</span> <span class="s">&quot;add&quot;</span><span class="p">)</span> <span class="p">{</span>
    <span class="n">Add</span><span class="p">(</span><span class="n">values</span><span class="p">.</span><span class="n">data</span><span class="p">(),</span> <span class="n">type_codes</span><span class="p">.</span><span class="n">data</span><span class="p">(),</span> <span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">());</span>
  <span class="p">}</span> <span class="k">else</span> <span class="k">if</span> <span class="p">(</span><span class="n">op_id_</span><span class="p">[</span><span class="n">id</span><span class="p">]</span> <span class="o">==</span> <span class="s">&quot;sub&quot;</span><span class="p">)</span> <span class="p">{</span>
    <span class="n">Sub</span><span class="p">(</span><span class="n">values</span><span class="p">.</span><span class="n">data</span><span class="p">(),</span> <span class="n">type_codes</span><span class="p">.</span><span class="n">data</span><span class="p">(),</span> <span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">());</span>
  <span class="p">}</span> <span class="k">else</span> <span class="k">if</span> <span class="p">(</span><span class="n">op_id_</span><span class="p">[</span><span class="n">id</span><span class="p">]</span> <span class="o">==</span> <span class="s">&quot;mul&quot;</span><span class="p">)</span> <span class="p">{</span>
    <span class="n">Mul</span><span class="p">(</span><span class="n">values</span><span class="p">.</span><span class="n">data</span><span class="p">(),</span> <span class="n">type_codes</span><span class="p">.</span><span class="n">data</span><span class="p">(),</span> <span class="n">args</span><span class="p">.</span><span class="n">size</span><span class="p">());</span>
  <span class="p">}</span> <span class="k">else</span> <span class="p">{</span>
    <span class="n">LOG</span><span class="p">(</span><span class="n">FATAL</span><span class="p">)</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;Unknown op: &quot;</span> <span class="o">&lt;&lt;</span> <span class="n">op_id_</span><span class="p">[</span><span class="n">id</span><span class="p">]</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
  <span class="p">}</span>
<span class="p">}</span>
</pre></div>
</div>
<p><code class="docutils literal notranslate"><span class="pre">Run</span></code> function mainly has two parts. The first part allocates a list of <code class="docutils literal notranslate"><span class="pre">TVMValue</span></code>, and maps corresponding data entry blocks. This will become the arguments of our operator functions. The second part than invokes our operator functions. Although we use the same C functions as the previous example, you can replace <code class="docutils literal notranslate"><span class="pre">Add</span></code>, <code class="docutils literal notranslate"><span class="pre">Sub</span></code>, and <code class="docutils literal notranslate"><span class="pre">Mul</span></code> with your own engine. You only need to make sure your engine stores the results to the last argument so that they can be transferred back to TVM runtime.</p>
<p>With above functions implemented, our customized codegen and runtime can now execute subgraphs. The last step is registering an API (<code class="docutils literal notranslate"><span class="pre">examplejson_module_create</span></code>) to create this module:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">TVM_REGISTER_GLOBAL</span><span class="p">(</span><span class="s">&quot;module.examplejson_module_create&quot;</span><span class="p">)</span>
<span class="p">.</span><span class="n">set_body_typed</span><span class="p">([](</span><span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">code</span><span class="p">){</span>
    <span class="k">auto</span> <span class="n">n</span> <span class="o">=</span> <span class="n">make_object</span><span class="o">&lt;</span><span class="n">ExampleJsonModule</span><span class="o">&gt;</span><span class="p">(</span><span class="n">code</span><span class="p">);</span>
    <span class="k">return</span> <span class="n">runtime</span><span class="o">::</span><span class="n">Module</span><span class="p">(</span><span class="n">n</span><span class="p">);</span>
<span class="p">});</span>
</pre></div>
</div>
</div>
<div class="section" id="implement-savetobinary-and-loadfrombinary">
<h4>Implement SaveToBinary and LoadFromBinary<a class="headerlink" href="#implement-savetobinary-and-loadfrombinary" title="永久链接至标题">¶</a></h4>
<p>So far we have implemented the main features of a customized runtime so that it can be used as other TVM runtimes. However, when users want to save the built runtime to a disk for deployment, TVM has no idea about how to save it. This is the reason we want to implement <code class="docutils literal notranslate"><span class="pre">SaveToBinary</span></code> and <code class="docutils literal notranslate"><span class="pre">LoadFromBinary</span></code>, which tell TVM how should this customized runtime be persist and restored.</p>
<p>We first implement <code class="docutils literal notranslate"><span class="pre">SaveToBinary</span></code> function to allow users to save this module in disk.</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="kt">void</span> <span class="nf">SaveToBinary</span><span class="p">(</span><span class="n">dmlc</span><span class="o">::</span><span class="n">Stream</span><span class="o">*</span> <span class="n">stream</span><span class="p">)</span> <span class="k">final</span> <span class="p">{</span>
    <span class="n">stream</span><span class="o">-&gt;</span><span class="n">Write</span><span class="p">(</span><span class="k">this</span><span class="o">-&gt;</span><span class="n">graph_json_</span><span class="p">);</span>
<span class="p">}</span>
</pre></div>
</div>
<p>We can find that this function is pretty simple. Recall that the only argument we took in constructor is a subgraph representation, meaning that we only need a subgraph representation to construct/recover this customized runtime module. As a result, <code class="docutils literal notranslate"><span class="pre">SaveToBinary</span></code> simply writes the subgraph to an output DMLC stream. That is, when users use <code class="docutils literal notranslate"><span class="pre">export_library</span></code> API to export the module, the customized module will be an ExampleJSON stream of a subgraph.</p>
<p>Similarity, <code class="docutils literal notranslate"><span class="pre">LoadFromBinary</span></code> reads the subgraph stream and re-constructs the customized runtime module:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="k">static</span> <span class="n">Module</span> <span class="nf">LoadFromBinary</span><span class="p">(</span><span class="kt">void</span><span class="o">*</span> <span class="n">strm</span><span class="p">)</span> <span class="p">{</span>
  <span class="n">dmlc</span><span class="o">::</span><span class="n">Stream</span><span class="o">*</span> <span class="n">stream</span> <span class="o">=</span> <span class="k">static_cast</span><span class="o">&lt;</span><span class="n">dmlc</span><span class="o">::</span><span class="n">Stream</span><span class="o">*&gt;</span><span class="p">(</span><span class="n">strm</span><span class="p">);</span>
  <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">graph_json</span><span class="p">;</span>
  <span class="n">stream</span><span class="o">-&gt;</span><span class="n">Read</span><span class="p">(</span><span class="o">&amp;</span><span class="n">graph_json</span><span class="p">);</span>
  <span class="k">auto</span> <span class="n">n</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">::</span><span class="n">runtime</span><span class="o">::</span><span class="n">make_object</span><span class="o">&lt;</span><span class="n">ExampleJsonModule</span><span class="o">&gt;</span><span class="p">(</span><span class="n">graph_json</span><span class="p">);</span>
  <span class="k">return</span> <span class="n">Module</span><span class="p">(</span><span class="n">n</span><span class="p">);</span>
<span class="p">}</span>
</pre></div>
</div>
<p>We also need to register this function to enable the corresponding Python API:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">TVM_REGISTER_GLOBAL</span><span class="p">(</span><span class="s">&quot;module.loadbinary_examplejson&quot;</span><span class="p">)</span>
<span class="p">.</span><span class="n">set_body_typed</span><span class="p">(</span><span class="n">ExampleJsonModule</span><span class="o">::</span><span class="n">LoadFromBinary</span><span class="p">);</span>
</pre></div>
</div>
<p>The above registration means when users call <code class="docutils literal notranslate"><span class="pre">tvm.runtime.load_module(lib_path)</span></code> API and the exported library has an ExampleJSON stream, our <code class="docutils literal notranslate"><span class="pre">LoadFromBinary</span></code> will be invoked to create the same customized runtime module.</p>
<p>In addition, if you want to support module creation directly from an ExampleJSON file, you can also implement a simple function and register a Python API as follows:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="k">static</span> <span class="n">Module</span> <span class="nf">Create</span><span class="p">(</span><span class="k">const</span> <span class="n">std</span><span class="o">::</span><span class="n">string</span><span class="o">&amp;</span> <span class="n">path</span><span class="p">)</span> <span class="p">{</span>
    <span class="n">std</span><span class="o">::</span><span class="n">ifstream</span> <span class="n">filep</span><span class="p">;</span>
    <span class="n">filep</span><span class="p">.</span><span class="n">open</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="n">std</span><span class="o">::</span><span class="n">ios</span><span class="o">::</span><span class="n">in</span><span class="p">);</span>
    <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">graph_json</span><span class="p">;</span>
    <span class="n">std</span><span class="o">::</span><span class="n">string</span> <span class="n">line</span><span class="p">;</span>
    <span class="k">while</span> <span class="p">(</span><span class="n">std</span><span class="o">::</span><span class="n">getline</span><span class="p">(</span><span class="n">filep</span><span class="p">,</span> <span class="n">line</span><span class="p">))</span> <span class="p">{</span>
        <span class="n">graph_json</span> <span class="o">+=</span> <span class="n">line</span><span class="p">;</span>
        <span class="n">graph_json</span> <span class="o">+=</span> <span class="s">&quot;</span><span class="se">\n</span><span class="s">&quot;</span><span class="p">;</span>
    <span class="p">}</span>
    <span class="n">filep</span><span class="p">.</span><span class="n">close</span><span class="p">();</span>
    <span class="k">auto</span> <span class="n">n</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">::</span><span class="n">runtime</span><span class="o">::</span><span class="n">make_object</span><span class="o">&lt;</span><span class="n">ExampleJsonModule</span><span class="o">&gt;</span><span class="p">(</span><span class="n">graph_json</span><span class="p">);</span>
    <span class="k">return</span> <span class="n">Module</span><span class="p">(</span><span class="n">n</span><span class="p">);</span>
<span class="p">}</span>

<span class="n">TVM_REGISTER_GLOBAL</span><span class="p">(</span><span class="s">&quot;module.loadfile_examplejson&quot;</span><span class="p">)</span>
<span class="p">.</span><span class="n">set_body</span><span class="p">([](</span><span class="n">TVMArgs</span> <span class="n">args</span><span class="p">,</span> <span class="n">TVMRetValue</span><span class="o">*</span> <span class="n">rv</span><span class="p">)</span> <span class="p">{</span>
    <span class="o">*</span><span class="n">rv</span> <span class="o">=</span> <span class="n">ExampleJsonModule</span><span class="o">::</span><span class="n">Create</span><span class="p">(</span><span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]);</span>
<span class="p">});</span>
</pre></div>
</div>
<p>It means users can manually write/modify an ExampleJSON file, and use Python API <code class="docutils literal notranslate"><span class="pre">tvm.runtime.load_module(&quot;mysubgraph.examplejson&quot;,</span> <span class="pre">&quot;examplejson&quot;)</span></code> to construct a customized module.</p>
</div>
</div>
</div>
<div class="section" id="summary">
<h2>总结<a class="headerlink" href="#summary" title="永久链接至标题">¶</a></h2>
<p>In summary, here is a checklist for you to refer:</p>
<ul class="simple">
<li><p>A codegen class derived from <code class="docutils literal notranslate"><span class="pre">ExprVisitor</span></code> and <code class="docutils literal notranslate"><span class="pre">CodegenCBase</span></code> (only for C codegen) with following functions.</p>
<ul>
<li><p><code class="docutils literal notranslate"><span class="pre">VisitExpr_(const</span> <span class="pre">CallNode*</span> <span class="pre">call)</span></code> to collect call node information.</p></li>
<li><p>Other visitor functions you needed to collect subgraph information.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">JIT</span></code> to generate subgraph code.</p></li>
<li><p>Register codegen.</p></li>
</ul>
</li>
<li><p>A function to create <code class="docutils literal notranslate"><span class="pre">CSourceModule</span></code> (for C codegen).</p></li>
<li><p>A runtime module class derived from <code class="docutils literal notranslate"><span class="pre">ModuleNode</span></code> with following functions (for your graph representation).</p>
<ul>
<li><p>Constructor.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">GetFunction</span></code> to generate a TVM runtime compatible <code class="docutils literal notranslate"><span class="pre">PackedFunc</span></code>.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">Run</span></code> to execute a subgraph.</p></li>
<li><p>Register a runtime creation API.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">SaveToBinary</span></code> and <code class="docutils literal notranslate"><span class="pre">LoadFromBinary</span></code> to serialize/deserialize customized runtime module.</p></li>
<li><p>Register <code class="docutils literal notranslate"><span class="pre">LoadFromBinary</span></code> API to support <code class="docutils literal notranslate"><span class="pre">tvm.runtime.load_module(your_module_lib_path)</span></code>.</p></li>
<li><p>(optional) <code class="docutils literal notranslate"><span class="pre">Create</span></code> to support customized runtime module construction from subgraph file in your representation.</p></li>
</ul>
</li>
<li><p>An annotator to annotate a user Relay program to make use of your compiler and runtime (TBA).</p></li>
</ul>
</div>
</div>


           </div>
           
          </div>
          

<footer>

    <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
      
        <a href="pytest_target_parametrization.html" class="btn btn-neutral float-right" title="Python Target Parametrization" accesskey="n" rel="next">下一个 <span class="fa fa-arrow-circle-right"></span></a>
      
      
        <a href="relay_add_pass.html" class="btn btn-neutral float-left" title="Adding a Compiler Pass to Relay" accesskey="p" rel="prev"><span class="fa fa-arrow-circle-left"></span> 上一个</a>
      
    </div>

<div id="button" class="backtop"><img src="../../_static//img/right.svg" alt="backtop"/> </div>
<section class="footerSec">
    <div class="footerHeader">
      <ul class="d-flex align-md-items-center justify-content-between flex-column flex-md-row">
        <li class="copywrite d-flex align-items-center">
          <h5 id="copy-right-info">© 2020 Apache Software Foundation | All right reserved</h5>
        </li>
      </ul>

    </div>

    <ul>
      <li class="footernote">Copyright © 2020 The Apache Software Foundation. Apache TVM, Apache, the Apache feather, and the Apache TVM project logo are either trademarks or registered trademarks of the Apache Software Foundation.</li>
    </ul>

</section>
</footer>
        </div>
      </div>

    </section>

  </div>
  

    <script src="https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.12.9/umd/popper.min.js" integrity="sha384-ApNbgh9B+Y1QKtv3Rn7W3mgPxhU9K/ScQsAP7hUibX39j7fakFPskvXusvfa0b4Q" crossorigin="anonymous"></script>
    <script src="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/js/bootstrap.min.js" integrity="sha384-JZR6Spejh4U02d8jOt6vLEHfe/JQGiRRSQQxSfFWpi1MquVdAyjUar5+76PVCmYl" crossorigin="anonymous"></script>

  </body>
  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
    <!-- Theme Analytics -->
    <script>
    (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){
      (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o),
      m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m)
    })(window,document,'script','https://www.google-analytics.com/analytics.js','ga');

    ga('create', 'UA-75982049-2', 'auto');
    ga('send', 'pageview');
    </script>

    
   

</body>
</html>